Adding a DatasetFolder class. (#444)

* Adding tests to ImageFolder

* Adding DatasetFolder class

* Fix tests for pytest and code for lint checker

* Adding mock to requirements for ImageFolder tests

* Remove mocks from requirements

Author: freud14

docs/source/datasets.rst

@@ -4,11 +4,11 @@ torchvision.datasets
 All datasets are subclasses of :class:`torch.utils.data.Dataset`
 i.e, they have ``__getitem__`` and ``__len__`` methods implemented.
 Hence, they can all be passed to a :class:`torch.utils.data.DataLoader`
-which can load multiple samples parallelly using ``torch.multiprocessing`` workers. 
+which can load multiple samples parallelly using ``torch.multiprocessing`` workers.
 For example: ::
-    
+
     imagenet_data = torchvision.datasets.ImageFolder('path/to/imagenet_root/')
-    data_loader = torch.utils.data.DataLoader(imagenet_data, 
+    data_loader = torch.utils.data.DataLoader(imagenet_data,
                                               batch_size=4,
                                               shuffle=True,
                                               num_workers=args.nThreads)
@@ -22,7 +22,7 @@ All the datasets have almost similar API. They all have two common arguments:
 ``transform`` and  ``target_transform`` to transform the input and target respectively.
 
 
-.. currentmodule:: torchvision.datasets 
+.. currentmodule:: torchvision.datasets
 
 
 MNIST
@@ -78,6 +78,14 @@ ImageFolder
   :members: __getitem__
   :special-members:
 
+DatasetFolder
+~~~~~~~~~~~~~
+
+.. autoclass:: DatasetFolder
+  :members: __getitem__
+  :special-members:
+
+
 
 Imagenet-12
 ~~~~~~~~~~~
@@ -121,4 +129,3 @@ PhotoTour
 .. autoclass:: PhotoTour
   :members: __getitem__
   :special-members:
-

test/assets/dataset/a/a1.png

@@ -0,0 +1,66 @@
+import unittest
+
+import os
+
+from torchvision.datasets import ImageFolder
+
+
+def mock_transform(return_value, arg_list):
+    def mock(arg):
+        arg_list.append(arg)
+        return return_value
+    return mock
+
+
+class Tester(unittest.TestCase):
+    root = 'test/assets/dataset/'
+    classes = ['a', 'b']
+    class_a_images = [os.path.join('test/assets/dataset/a/', path) for path in ['a1.png', 'a2.png', 'a3.png']]
+    class_b_images = [os.path.join('test/assets/dataset/b/', path) for path in ['b1.png', 'b2.png', 'b3.png', 'b4.png']]
+
+    def test_image_folder(self):
+        dataset = ImageFolder(Tester.root, loader=lambda x: x)
+        self.assertEqual(sorted(Tester.classes), sorted(dataset.classes))
+        for cls in Tester.classes:
+            self.assertEqual(cls, dataset.classes[dataset.class_to_idx[cls]])
+        class_a_idx = dataset.class_to_idx['a']
+        class_b_idx = dataset.class_to_idx['b']
+        imgs_a = [(img_path, class_a_idx)for img_path in Tester.class_a_images]
+        imgs_b = [(img_path, class_b_idx)for img_path in Tester.class_b_images]
+        imgs = sorted(imgs_a + imgs_b)
+        self.assertEqual(imgs, dataset.imgs)
+
+        outputs = sorted([dataset[i] for i in range(len(dataset))])
+        self.assertEqual(imgs, outputs)
+
+    def test_transform(self):
+        return_value = 'test/assets/dataset/a/a1.png'
+
+        args = []
+        transform = mock_transform(return_value, args)
+
+        dataset = ImageFolder(Tester.root, loader=lambda x: x, transform=transform)
+        outputs = [dataset[i][0] for i in range(len(dataset))]
+        self.assertEqual([return_value] * len(outputs), outputs)
+
+        imgs = sorted(Tester.class_a_images + Tester.class_b_images)
+        self.assertEqual(imgs, sorted(args))
+
+    def test_target_transform(self):
+        return_value = 1
+
+        args = []
+        target_transform = mock_transform(return_value, args)
+
+        dataset = ImageFolder(Tester.root, loader=lambda x: x, target_transform=target_transform)
+        outputs = [dataset[i][1] for i in range(len(dataset))]
+        self.assertEqual([return_value] * len(outputs), outputs)
+
+        class_a_idx = dataset.class_to_idx['a']
+        class_b_idx = dataset.class_to_idx['b']
+        targets = sorted([class_a_idx] * len(Tester.class_a_images) +
+                         [class_b_idx] * len(Tester.class_b_images))
+        self.assertEqual(targets, sorted(args))
+
+if __name__ == '__main__':
+    unittest.main()

test/assets/dataset/a/a2.png

@@ -1,5 +1,5 @@
 from .lsun import LSUN, LSUNClass
-from .folder import ImageFolder
+from .folder import ImageFolder, DatasetFolder
 from .coco import CocoCaptions, CocoDetection
 from .cifar import CIFAR10, CIFAR100
 from .stl10 import STL10
@@ -11,7 +11,7 @@
 from .omniglot import Omniglot
 
 __all__ = ('LSUN', 'LSUNClass',
-           'ImageFolder', 'FakeData',
+           'ImageFolder', 'DatasetFolder', 'FakeData',
            'CocoCaptions', 'CocoDetection',
            'CIFAR10', 'CIFAR100', 'EMNIST', 'FashionMNIST',
            'MNIST', 'STL10', 'SVHN', 'PhotoTour', 'SEMEION',