Adding a DatasetFolder class. (#442)

* Adding tests to ImageFolder

* Adding DatasetFolder class

* Fix tests for pytest and code for lint checker

* Adding mock to requirements for ImageFolder tests

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:
-

setup.py

@@ -33,6 +33,7 @@ def find_version(*file_paths):
     'pillow >= 4.1.1',
     'six',
     'torch',
+    'mock',
 ]
 
 setup(

test/assets/dataset/a/a1.png

@@ -0,0 +1,59 @@
+import unittest
+try:
+    from unittest.mock import Mock
+except ImportError as e:
+    from mock import Mock
+
+import os
+
+from torchvision.datasets import ImageFolder
+
+
+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'
+        transform = Mock(return_value=return_value)
+        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)
+        args = [call[0][0] for call in transform.call_args_list]
+        self.assertEqual(imgs, sorted(args))
+
+    def test_target_transform(self):
+        return_value = 1
+        target_transform = Mock(return_value=return_value)
+        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))
+        args = [call[0][0] for call in target_transform.call_args_list]
+        self.assertEqual(targets, sorted(args))
+
+if __name__ == '__main__':
+    unittest.main()