pytorch 源码片段：torchvision.datasets.ImageFolder

Source code for torchvision.datasets.folder

def find_classes(dir):
    classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
    classes.sort()
    class_to_idx = {classes[i]: i for i in range(len(classes))}
    return classes, class_to_idx

def make_dataset(dir, class_to_idx):
    images = []
    dir = os.path.expanduser(dir)
    for target in sorted(os.listdir(dir)):
        d = os.path.join(dir, target)
        if not os.path.isdir(d):
            continue

        for root, _, fnames in sorted(os.walk(d)):
            for fname in sorted(fnames):
                if is_image_file(fname):
                    path = os.path.join(root, fname)
                    item = (path, class_to_idx[target])
                    images.append(item)

    return images

class ImageFolder(data.Dataset):
    def __init__(self, root, transform=None, target_transform=None,
                 loader=default_loader):
        classes, class_to_idx = find_classes(root)
        imgs = make_dataset(root, class_to_idx)
        .
        .
        .
    def __getitem__(self, index):
        path, target = self.imgs[index]
        img = self.loader(path)
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

解读 find_classes(dir):
test 文件夹下存在 1，2，13，4，05 这五个空文件夹
os.listdir() 用于返回指定的文件夹包含的文件或文件夹的名字的列表
os.path.isdir() 用于判断某一路径是否为目录

import os

dir='C:/Users/Dell/Desktop/test'
classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))] #列表
classes.sort()
class_to_idx = {classes[i]:i for i in [4,1,2,0,3]} #字典
print('classes: ', classes)
print('class_to_idx: ',class_to_idx)

输出结果：
classes: ['05', '1', '13', '2', '4'] 
class_to_idx: {'4': 4, '1': 1, '13': 2, '05': 0, '2': 3}

classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]

等价于：

classes = []
for d in os.listdir(dir):
    if os.path.isdir(os.path.join(dir, d)):
        classes.append(d)