A generic data loader where the images are arranged in this way: :: root/dog/xxx.png root/dog/xxy.png root/dog/xxz.png root/cat/123.png root/cat/nsdf3.png root/cat/asd932_.png Args: root (string): Root directory path. transform (
| 76 | |
| 77 | |
| 78 | class ImageFolder(data.Dataset): |
| 79 | """A generic data loader where the images are arranged in this way: :: |
| 80 | |
| 81 | root/dog/xxx.png |
| 82 | root/dog/xxy.png |
| 83 | root/dog/xxz.png |
| 84 | |
| 85 | root/cat/123.png |
| 86 | root/cat/nsdf3.png |
| 87 | root/cat/asd932_.png |
| 88 | |
| 89 | Args: |
| 90 | root (string): Root directory path. |
| 91 | transform (callable, optional): A function/transform that takes in an PIL image |
| 92 | and returns a transformed version. E.g, ``transforms.RandomCrop`` |
| 93 | target_transform (callable, optional): A function/transform that takes in the |
| 94 | target and transforms it. |
| 95 | loader (callable, optional): A function to load an image given its path. |
| 96 | |
| 97 | Attributes: |
| 98 | classes (list): List of the class names. |
| 99 | class_to_idx (dict): Dict with items (class_name, class_index). |
| 100 | imgs (list): List of (image path, class_index) tuples |
| 101 | """ |
| 102 | |
| 103 | def __init__(self, root, transform=None, target_transform=None, |
| 104 | loader=default_loader, classes_idx=None): |
| 105 | self.classes_idx = classes_idx |
| 106 | classes, class_to_idx = find_classes(root, self.classes_idx) |
| 107 | imgs = make_dataset(root, class_to_idx) |
| 108 | if len(imgs) == 0: |
| 109 | raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n" |
| 110 | "Supported image extensions are: " + ",".join(IMG_EXTENSIONS))) |
| 111 | |
| 112 | self.root = root |
| 113 | self.imgs = imgs |
| 114 | self.classes = classes |
| 115 | self.class_to_idx = class_to_idx |
| 116 | self.transform = transform |
| 117 | self.target_transform = target_transform |
| 118 | self.loader = loader |
| 119 | |
| 120 | def __getitem__(self, index): |
| 121 | """ |
| 122 | Args: |
| 123 | index (int): Index |
| 124 | |
| 125 | Returns: |
| 126 | tuple: (image, target) where target is class_index of the target class. |
| 127 | """ |
| 128 | path, target = self.imgs[index] |
| 129 | img = self.loader(path) |
| 130 | if self.transform is not None: |
| 131 | img = self.transform(img) |
| 132 | if self.target_transform is not None: |
| 133 | target = self.target_transform(target) |
| 134 | |
| 135 | return img, target |