Pytorch ImageFoader源码_imageformats.extensions-优快云博客

本文链接：https://blog.youkuaiyun.com/xuyunyunaixuexi/article/details/100698216

Pytorch的ImageFolder模块用于加载和组织指定根目录下的类目图像，通过transform和target_transform对图像和标签进行预处理。默认使用PIL加载图片，返回内容包括类别列表、类别到索引的映射以及图片路径与类别的元组。经过transform处理的图片将返回Tensor，否则返回PIL Image对象。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch.utils.data as data
#PIL: Python Image Library缩写，图像处理模块
#     Image,ImageFont,ImageDraw,ImageFilter
from PIL import Image    
import os
import os.path
 
# 图片扩展（图片格式）
IMG_EXTENSIONS = [
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
 
# 判断是不是图片文件
def is_image_file(filename):
    # 只要文件以IMG_EXTENSIONS结尾，就是图片
    # 注意any的使用
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
 
# 结果:classes:['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
# classes_to_idx:{'1': 1, '0': 0, '3': 3, '2': 2, '5': 5, '4': 4, '7': 7, '6': 6, '9': 9, '8': 8}
def find_classes(dir):
    '''
     返回dir下的类别名，classes:所有的类别，class_to_idx:将文件中str的类别名转化为int类别
     classes为目录下所有文件夹名字的集合
    '''
    # os.listdir：以列表的形式显示当前目录下的所有文件名和目录名，但不会区分文件和目录。
    # os.path.isdir：判定对象是否是目录，是则返回True，否则返回False
    # os.path.join：连接目录和文件名
 
    classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
    # sort:排序
    classes.sort()
    # 将文件名中得到的类别转化为数字class_to_idx['3'] = 3
    class_to_idx = {classes[i]: i for i in range(len(classes))}
    return classes, class_to_idx
    # class_to_idx :{