crashedice/signify

1"""A modified image folder class2 3We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)4so that this class can load images from both current directory and its subdirectories.5"""6 7import torch.utils.data as data8 9from PIL import Image10import os11 12IMG_EXTENSIONS = [13    '.jpg', '.JPG', '.jpeg', '.JPEG',14    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',15    '.tif', '.TIF', '.tiff', '.TIFF',16]17 18 19def is_image_file(filename):20    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)21 22 23def make_dataset(dir, max_dataset_size=float("inf")):24    images = []25    assert os.path.isdir(dir), '%s is not a valid directory' % dir26 27    for root, _, fnames in sorted(os.walk(dir)):28        for fname in fnames:29            if is_image_file(fname):30                path = os.path.join(root, fname)31                images.append(path)32    return images[:min(max_dataset_size, len(images))]33 34 35def default_loader(path):36    return Image.open(path).convert('RGB')37 38 39class ImageFolder(data.Dataset):40 41    def __init__(self, root, transform=None, return_paths=False,42                 loader=default_loader):43        imgs = make_dataset(root)44        if len(imgs) == 0:45            raise(RuntimeError("Found 0 images in: " + root + "\n"46                               "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))47 48        self.root = root49        self.imgs = imgs50        self.transform = transform51        self.return_paths = return_paths52        self.loader = loader53 54    def __getitem__(self, index):55        path = self.imgs[index]56        img = self.loader(path)57        if self.transform is not None:58            img = self.transform(img)59        if self.return_paths:60            return img, path61        else:62            return img63 64    def __len__(self):65        return len(self.imgs)66