Add ImageFolderWithoutTargets

torchvision/datasets/__init__.py

@@ -1,11 +1,8 @@
 from .lsun import LSUN, LSUNClass
-from .folder import ImageFolder
+from .folder import ImageClassFolder, ImageFolder
 from .coco import CocoCaptions, CocoDetection
 from .cifar import CIFAR10, CIFAR100
 from .mnist import MNIST
 
-__all__ = ('LSUN', 'LSUNClass',
-           'ImageFolder',
-           'CocoCaptions', 'CocoDetection',
-           'CIFAR10', 'CIFAR100',
-           'MNIST')
+__all__ = ('LSUN', 'LSUNClass', 'ImageFolder', 'ImageClassFolder',
+           'CocoCaptions', 'CocoDetection', 'CIFAR10', 'CIFAR100', 'MNIST')

torchvision/datasets/folder.py

@@ -1,23 +1,34 @@
-import torch.utils.data as data
-
-from PIL import Image
 import os
 import os.path
 
+import torch.utils.data as data
+from PIL import Image
+
 IMG_EXTENSIONS = [
-    '.jpg', '.JPG', '.jpeg', '.JPEG',
-    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+    '.jpg',
+    '.JPG',
+    '.jpeg',
+    '.JPEG',
+    '.png',
+    '.PNG',
+    '.ppm',
+    '.PPM',
+    '.bmp',
+    '.BMP',
 ]
 
+
 def is_image_file(filename):
     return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
 
+
 def find_classes(dir):
     classes = os.listdir(dir)
     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 = []
     for target in os.listdir(dir):
@@ -38,8 +49,11 @@ def default_loader(path):
     return Image.open(path).convert('RGB')
 
 
-class ImageFolder(data.Dataset):
-    def __init__(self, root, transform=None, target_transform=None,
+class ImageClassFolder(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)
@@ -64,3 +78,28 @@ def __getitem__(self, index):
 
     def __len__(self):
         return len(self.imgs)
+
+
+class ImageFolder(data.Dataset):
+    """ ImageFolder can be used to load images where there are no labels."""
+
+    def __init__(self, root, transform=None, loader=default_loader):
+        images = []
+        for filename in os.listdir(root):
+            if is_image_file(filename):
+                images.append('{}'.format(filename))
+
+        self.root = root
+        self.imgs = images
+        self.transform = transform
+        self.loader = loader
+
+    def __getitem__(self, index):
+        filename = self.imgs[index]
+        img = self.loader(os.path.join(self.root, filename))
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, filename
+
+    def __len__(self):
+        return len(self.imgs)