Skip to content

GitLab

Commit 7d8cc199 authored by Ross Taylor's avatar Ross Taylor Committed by Francisco Massa
Browse files

Add tests for Cityscapes Dataset (#1079) 

* Cityscapes test; semantic segmentation getitem test

* Add multiple targets test for Cityscapes

* flake8 fix for test_datasets
parent 6c560290
Hide whitespace changes
Inline Side-by-side
Showing with 131 additions and 5 deletions
test/fakedata_generation.py
View file @ 7d8cc199
......@@ -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
View file @ 7d8cc199
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
View file @ 7d8cc199
......@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment