Add tests for Cityscapes Dataset (#1079)

* Cityscapes test; semantic segmentation getitem test

* Add multiple targets test for Cityscapes

* flake8 fix for test_datasets

test/fakedata_generation.py

@@ -2,6 +2,7 @@ import os
 import sys
 import contextlib
 import tarfile
+import json
 import numpy as np
 import PIL
 import torch
@@ -168,3 +169,75 @@ def imagenet_root():
         _make_devkit_archive(root)
 
         yield root
+
+
+@contextlib.contextmanager
+def cityscapes_root():
+
+    def _make_image(file):
+        PIL.Image.fromarray(np.zeros((1024, 2048, 3), dtype=np.uint8)).save(file)
+
+    def _make_regular_target(file):
+        PIL.Image.fromarray(np.zeros((1024, 2048), dtype=np.uint8)).save(file)
+
+    def _make_color_target(file):
+        PIL.Image.fromarray(np.zeros((1024, 2048, 4), dtype=np.uint8)).save(file)
+
+    def _make_polygon_target(file):
+        polygon_example = {
+            'imgHeight': 1024,
+            'imgWidth': 2048,
+            'objects': [{'label': 'sky',
+                         'polygon': [[1241, 0], [1234, 156],
+                                     [1478, 197], [1611, 172],
+                                     [1606, 0]]},
+                        {'label': 'road',
+                         'polygon': [[0, 448], [1331, 274],
+                                     [1473, 265], [2047, 605],
+                                     [2047, 1023], [0, 1023]]}]}
+        with open(file, 'w') as outfile:
+            json.dump(polygon_example, outfile)
+
+    with get_tmp_dir() as tmp_dir:
+
+        for mode in ['Coarse', 'Fine']:
+            gt_dir = os.path.join(tmp_dir, 'gt%s' % mode)
+            os.makedirs(gt_dir)
+
+            if mode == 'Coarse':
+                splits = ['train', 'train_extra', 'val']
+            else:
+                splits = ['train', 'test', 'val']
+
+            for split in splits:
+                split_dir = os.path.join(gt_dir, split)
+                os.makedirs(split_dir)
+                for city in ['bochum', 'bremen']:
+                    city_dir = os.path.join(split_dir, city)
+                    os.makedirs(city_dir)
+                    _make_color_target(os.path.join(city_dir,
+                                                    '{city}_000000_000000_gt{mode}_color.png'.format(
+                                                        city=city, mode=mode)))
+                    _make_regular_target(os.path.join(city_dir,
+                                                      '{city}_000000_000000_gt{mode}_instanceIds.png'.format(
+                                                          city=city, mode=mode)))
+                    _make_regular_target(os.path.join(city_dir,
+                                                      '{city}_000000_000000_gt{mode}_labelIds.png'.format(
+                                                          city=city, mode=mode)))
+                    _make_polygon_target(os.path.join(city_dir,
+                                                      '{city}_000000_000000_gt{mode}_polygons.json'.format(
+                                                          city=city, mode=mode)))
+
+        # leftImg8bit dataset
+        leftimg_dir = os.path.join(tmp_dir, 'leftImg8bit')
+        os.makedirs(leftimg_dir)
+        for split in ['test', 'train_extra', 'train', 'val']:
+            split_dir = os.path.join(leftimg_dir, split)
+            os.makedirs(split_dir)
+            for city in ['bochum', 'bremen']:
+                city_dir = os.path.join(split_dir, city)
+                os.makedirs(city_dir)
+                _make_image(os.path.join(city_dir,
+                                         '{city}_000000_000000_leftImg8bit.png'.format(city=city)))
+
+        yield tmp_dir

test/test_datasets.py

 import os
 import unittest
 import mock
+import numpy as np
 import PIL
 from PIL import Image
 from torch._utils_internal import get_file_path_2
 import torchvision
 from common_utils import get_tmp_dir
-from fakedata_generation import mnist_root, cifar_root, imagenet_root
+from fakedata_generation import mnist_root, cifar_root, imagenet_root, cityscapes_root
 
 
 class Tester(unittest.TestCase):
@@ -16,6 +17,12 @@ class Tester(unittest.TestCase):
         self.assertTrue(isinstance(img, PIL.Image.Image))
         self.assertTrue(isinstance(target, int))
 
+    def generic_segmentation_dataset_test(self, dataset, num_images=1):
+        self.assertEqual(len(dataset), num_images)
+        img, target = dataset[0]
+        self.assertTrue(isinstance(img, PIL.Image.Image))
+        self.assertTrue(isinstance(target, PIL.Image.Image))
+
     def test_imagefolder(self):
         # TODO: create the fake data on-the-fly
         FAKEDATA_DIR = get_file_path_2(
@@ -133,6 +140,51 @@ class Tester(unittest.TestCase):
             img, target = dataset[0]
             self.assertEqual(dataset.class_to_idx[dataset.classes[0]], target)
 
+    def test_cityscapes(self):
+        with cityscapes_root() as root:
+
+            for mode in ['coarse', 'fine']:
+
+                if mode == 'coarse':
+                    splits = ['train', 'train_extra', 'val']
+                else:
+                    splits = ['train', 'val', 'test']
+
+                for split in splits:
+                    for target_type in ['semantic', 'instance']:
+                        dataset = torchvision.datasets.Cityscapes(root, split=split,
+                                                                  target_type=target_type, mode=mode)
+                        self.generic_segmentation_dataset_test(dataset, num_images=2)
+
+                    color_dataset = torchvision.datasets.Cityscapes(root, split=split,
+                                                                    target_type='color', mode=mode)
+                    color_img, color_target = color_dataset[0]
+                    self.assertTrue(isinstance(color_img, PIL.Image.Image))
+                    self.assertTrue(np.array(color_target).shape[2] == 4)
+
+                    polygon_dataset = torchvision.datasets.Cityscapes(root, split=split,
+                                                                      target_type='polygon', mode=mode)
+                    polygon_img, polygon_target = polygon_dataset[0]
+                    self.assertTrue(isinstance(polygon_img, PIL.Image.Image))
+                    self.assertTrue(isinstance(polygon_target, dict))
+                    self.assertTrue(isinstance(polygon_target['imgHeight'], int))
+                    self.assertTrue(isinstance(polygon_target['objects'], list))
+
+                    # Test multiple target types
+                    targets_combo = ['semantic', 'polygon', 'color']
+                    multiple_types_dataset = torchvision.datasets.Cityscapes(root, split=split,
+                                                                             target_type=targets_combo,
+                                                                             mode=mode)
+                    output = multiple_types_dataset[0]
+                    self.assertTrue(isinstance(output, tuple))
+                    self.assertTrue(len(output) == 2)
+                    self.assertTrue(isinstance(output[0], PIL.Image.Image))
+                    self.assertTrue(isinstance(output[1], tuple))
+                    self.assertTrue(len(output[1]) == 3)
+                    self.assertTrue(isinstance(output[1][0], PIL.Image.Image))  # semantic
+                    self.assertTrue(isinstance(output[1][1], dict))  # polygon
+                    self.assertTrue(isinstance(output[1][2], PIL.Image.Image))  # color
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/datasets/cityscapes.py

@@ -127,15 +127,16 @@ class Cityscapes(VisionDataset):
                              ' or "color"')
 
         if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir):
+
             if split == 'train_extra':
-                image_dir_zip = os.path.join(self.root, 'leftImg8bit') + '_trainextra.zip'
+                image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainextra.zip'))
             else:
-                image_dir_zip = os.path.join(self.root, 'leftImg8bit') + '_trainvaltest.zip'
+                image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainvaltest.zip'))
 
             if self.mode == 'gtFine':
-                target_dir_zip = os.path.join(self.root, self.mode) + '_trainvaltest.zip'
+                target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '_trainvaltest.zip'))
             elif self.mode == 'gtCoarse':
-                target_dir_zip = os.path.join(self.root, self.mode)
+                target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '.zip'))
 
             if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip):
                 extract_archive(from_path=image_dir_zip, to_path=self.root)