在使用pytorch构建数据库时,会使用到ImageFolder这个模块便于数据加载,了解其源码便于快速开发。
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':