diff --git a/tensorflow_datasets/image/__init__.py b/tensorflow_datasets/image/__init__.py index 599e5b625a5..cc7a678edc7 100644 --- a/tensorflow_datasets/image/__init__.py +++ b/tensorflow_datasets/image/__init__.py @@ -34,6 +34,7 @@ from tensorflow_datasets.image.cifar10_1 import Cifar10_1 from tensorflow_datasets.image.cifar10_corrupted import Cifar10Corrupted from tensorflow_datasets.image.citrus import CitrusLeaves +from tensorflow_datasets.image.cityscapes import Cityscapes from tensorflow_datasets.image.clevr import CLEVR from tensorflow_datasets.image.cmaterdb import Cmaterdb from tensorflow_datasets.image.coil100 import Coil100 diff --git a/tensorflow_datasets/image/cityscapes.py b/tensorflow_datasets/image/cityscapes.py new file mode 100644 index 00000000000..8df1ff69639 --- /dev/null +++ b/tensorflow_datasets/image/cityscapes.py @@ -0,0 +1,282 @@ +'''Cityscapes Datasets.''' + +import os +import re + +import tensorflow as tf +import tensorflow_datasets.public_api as tfds +from tensorflow_datasets.core import api_utils + +_CITATION = '''\ +@inproceedings{Cordts2016Cityscapes, + title={The Cityscapes Dataset for Semantic Urban Scene Understanding}, + author={Cordts, Marius and Omran, Mohamed and Ramos, Sebastian and Rehfeld, Timo and Enzweiler, Markus and Benenson, Rodrigo and Franke, Uwe and Roth, Stefan and Schiele, Bernt}, + booktitle={Proc. of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, + year={2016} +} +''' + +_DESCRIPTION = '''\ + Cityscapes is a dataset consisting of diverse urban street scenes across 50 different cities + at varying times of the year as well as ground truths for several vision tasks including + semantic segmentation, instance level segmentation (TODO), and stereo pair disparity inference. + + + For segmentation tasks (default split, accessible via 'cityscapes/semantic_segmentation'), Cityscapes provides + dense pixel level annotations for 5000 images at 1024 * 2048 resolution pre-split into training (2975), + validation (500) and test (1525) sets. Label annotations for segmentation tasks span across 30+ classes + commonly encountered during driving scene perception. Detailed label information may be found here: + https://github.com/mcordts/cityscapesScripts/blob/master/cityscapesscripts/helpers/labels.py#L52-L99 + + Cityscapes also provides coarse grain segmentation annotations (accessible via 'cityscapes/semantic_segmentation_extra') + for 19998 images in a 'train_extra' split which may prove useful for pretraining / data-heavy models. + + + Besides segmentation, cityscapes also provides stereo image pairs and ground truths for disparity inference + tasks on both the normal and extra splits (accessible via 'cityscapes/stereo_disparity' and + 'cityscapes/stereo_disparity_extra' respectively). + + Ingored examples: + - For 'cityscapes/stereo_disparity_extra': + - troisdorf_000000_000073_{*} images (no disparity map present) + + WARNING: this dataset requires users to setup a login and password in order to get the files. +''' + +# TODO add instance ids (might need to import cityScapesScripts) + +class CityscapesConfig(tfds.core.BuilderConfig): + '''BuilderConfig for Cityscapes + + Args: + right_images (bool): Enables right images for stereo image tasks. + segmentation_labels (bool): Enables image segmentation labels. + disparity_maps (bool): Enables disparity maps. + train_extra_split (bool): Enables train_extra split. This automatically + enables coarse grain segmentations, if segmentation labels are used. + ''' + + @api_utils.disallow_positional_args + def __init__(self, right_images=False, segmentation_labels=True, + disparity_maps=False, train_extra_split=False, **kwargs): + super().__init__(**kwargs) + + self.right_images = right_images + self.segmentation_labels = segmentation_labels + self.disparity_maps = disparity_maps + self.train_extra_split = train_extra_split + + self.ignored_ids = set() + + # Setup required zips and their root dir names + self.zip_root = {} + self.zip_root['images_left'] = ( + 'leftImg8bit_trainvaltest.zip', 'leftImg8bit') + + if self.train_extra_split: + self.zip_root['images_left/extra'] = ( + 'leftImg8bit_trainextra.zip', 'leftImg8bit') + + if self.right_images: + self.zip_root['images_right'] = ( + 'rightImg8bit_trainvaltest.zip', 'rightImg8bit') + if self.train_extra_split: + self.zip_root['images_right/extra'] = ( + 'rightImg8bit_trainextra.zip', 'rightImg8bit') + + if self.segmentation_labels: + if not self.train_extra_split: + self.zip_root['segmentation_labels'] = ( + 'gtFine_trainvaltest.zip', 'gtFine') + self.label_suffix = 'gtFine_labelIds' + else: + # The 'train extra' split only has coarse labels unlike train and val. + # Therefore, for consistency across splits, we also enable coarse labels + # using the train_extra_split flag. + self.zip_root['segmentation_labels'] = ('gtCoarse.zip', 'gtCoarse') + self.zip_root['segmentation_labels/extra'] = ( + 'gtCoarse.zip', 'gtCoarse') + self.label_suffix = 'gtCoarse_labelIds' + + if self.disparity_maps: + self.zip_root['disparity_maps'] = ( + 'disparity_trainvaltest.zip', 'disparity') + if self.train_extra_split: + self.zip_root['disparity_maps/extra'] = ( + 'disparity_trainextra.zip', 'disparity') + self.ignored_ids.add('troisdorf_000000_000073') # No disparity for this file + + +class Cityscapes(tfds.core.GeneratorBasedBuilder): + '''Base class for Cityscapes datasets''' + + BUILDER_CONFIGS = [ + CityscapesConfig( + name='semantic_segmentation', + description='Cityscapes semantic segmentation dataset.', + version="1.0.0", + right_images=False, + segmentation_labels=True, + disparity_maps=False, + train_extra_split=False, + ), + CityscapesConfig( + name='semantic_segmentation_extra', + description='Cityscapes semantic segmentation dataset with train_extra split and coarse labels.', # pylint: disable=line-too-long + version="1.0.0", + right_images=False, + segmentation_labels=True, + disparity_maps=False, + train_extra_split=True, + ), + CityscapesConfig( + name='stereo_disparity', + description='Cityscapes stereo image and disparity maps dataset.', + version="1.0.0", + right_images=True, + segmentation_labels=False, + disparity_maps=True, + train_extra_split=False, + ), + CityscapesConfig( + name='stereo_disparity_extra', + description='Cityscapes stereo image and disparity maps dataset with train_extra split.', # pylint: disable=line-too-long + version="1.0.0", + right_images=True, + segmentation_labels=False, + disparity_maps=True, + train_extra_split=True, + ), + ] + + VERSION = tfds.core.Version('1.0.0') + + def _info(self): + # Enable features as necessary + features = {} + features['image_id'] = tfds.features.Text() + features['image_left'] = tfds.features.Image( + shape=(1024, 2048, 3), encoding_format='png') + + if self.builder_config.right_images: + features['image_right'] = tfds.features.Image( + shape=(1024, 2048, 3), encoding_format='png') + + if self.builder_config.segmentation_labels: + features['segmentation_label'] = tfds.features.Image( + shape=(1024, 2048, 1), encoding_format='png') + + if self.builder_config.disparity_maps: + features['disparity_map'] = tfds.features.Image( + shape=(1024, 2048, 1), encoding_format='png') + + return tfds.core.DatasetInfo( + builder=self, + description=(_DESCRIPTION), + features=tfds.features.FeaturesDict(features), + homepage='https://www.cityscapes-dataset.com', + citation=_CITATION, + ) + + def _split_generators(self, dl_manager): + paths = {} + for split, (zip_file, zip_root) in self.builder_config.zip_root.items(): + paths[split] = os.path.join(dl_manager.manual_dir, zip_file) + + if any(not os.path.exists(z) for z in paths.values()): + msg = 'You must download the dataset files manually and place them in: ' + msg += ', '.join(paths.values()) + raise AssertionError(msg) + + for split, (_, zip_root) in self.builder_config.zip_root.items(): + paths[split] = os.path.join(dl_manager.extract(paths[split]), zip_root) + + splits = [ + tfds.core.SplitGenerator( + name=tfds.Split.TRAIN, + gen_kwargs={ + feat_dir: os.path.join(path, 'train') + for feat_dir, path in paths.items() + if not feat_dir.endswith('/extra') + }, + ), + tfds.core.SplitGenerator( + name=tfds.Split.VALIDATION, + gen_kwargs={ + feat_dir: os.path.join(path, 'val') + for feat_dir, path in paths.items() + if not feat_dir.endswith('/extra') + }, + ), + ] + + # Test split does not exist in coarse dataset + if not self.builder_config.train_extra_split: + splits.append(tfds.core.SplitGenerator( + name=tfds.Split.TEST, + gen_kwargs={ + feat_dir: os.path.join(path, 'test') + for feat_dir, path in paths.items() + if not feat_dir.endswith('/extra') + }, + )) + else: + splits.append(tfds.core.SplitGenerator( + name='train_extra', + gen_kwargs={ + feat_dir.replace('/extra', ''): os.path.join(path, 'train_extra') + for feat_dir, path in paths.items() + if feat_dir.endswith('/extra') + }, + )) + return splits + + def _generate_examples(self, **paths): + left_imgs_root = paths['images_left'] + for city_id in tf.io.gfile.listdir(left_imgs_root): + paths_city_root = {feat_dir: os.path.join(path, city_id) + for feat_dir, path in paths.items()} + + left_city_root = paths_city_root['images_left'] + for left_img in tf.io.gfile.listdir(left_city_root): + left_img_path = os.path.join(left_city_root, left_img) + image_id = _get_left_image_id(left_img) + + if image_id in self.builder_config.ignored_ids: + continue + + features = { + 'image_id': image_id, + 'image_left': left_img_path, + } + + if self.builder_config.right_images: + features['image_right'] = os.path.join( + paths_city_root['images_right'], + '{}_rightImg8bit.png'.format(image_id)) + + if self.builder_config.segmentation_labels: + features['segmentation_label'] = os.path.join( + paths_city_root['segmentation_labels'], + '{}_{}.png'.format( + image_id, self.builder_config.label_suffix)) + + if self.builder_config.disparity_maps: + features['disparity_map'] = os.path.join( + paths_city_root['disparity_maps'], + '{}_disparity.png'.format(image_id)) + + yield image_id, features + +# Helper functions + +LEFT_IMAGE_FILE_RE = re.compile(r'([a-z\-]+)_(\d+)_(\d+)_leftImg8bit\.png') + +def _get_left_image_id(left_image): + '''Returns the id of an image file. Used to associate an image file + with its corresponding label. + Example: + 'bonn_000001_000019_leftImg8bit' -> 'bonn_000001_000019' + ''' + match = LEFT_IMAGE_FILE_RE.match(left_image) + return '{}_{}_{}'.format(*match.groups()) diff --git a/tensorflow_datasets/image/cityscapes_test.py b/tensorflow_datasets/image/cityscapes_test.py new file mode 100644 index 00000000000..ea729bf2ff7 --- /dev/null +++ b/tensorflow_datasets/image/cityscapes_test.py @@ -0,0 +1,53 @@ + +'''Tests for Cityscapes dataset module.''' + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow_datasets import testing +from tensorflow_datasets.image import cityscapes + +# TODO add tests for features and files per configuration +class CityscapesSegmentationTest(testing.DatasetBuilderTestCase): + DATASET_CLASS = cityscapes.Cityscapes + BUILDER_CONFIG_NAMES_TO_TEST = ['semantic_segmentation'] + SPLITS = { + 'train': 3, + 'validation': 1, + 'test': 2, + } + + +class CityscapesSegmentationExtraTest(testing.DatasetBuilderTestCase): + DATASET_CLASS = cityscapes.Cityscapes + BUILDER_CONFIG_NAMES_TO_TEST = ['semantic_segmentation_extra'] + SPLITS = { + 'train': 3, + 'train_extra': 4, + 'validation': 1, + } + + +class CityscapesStereoDisparityTest(testing.DatasetBuilderTestCase): + DATASET_CLASS = cityscapes.Cityscapes + BUILDER_CONFIG_NAMES_TO_TEST = ['stereo_disparity'] + SPLITS = { + 'train': 3, + 'validation': 1, + 'test': 2, + } + + +class CityscapesStereoDisparityExtraTest(testing.DatasetBuilderTestCase): + DATASET_CLASS = cityscapes.Cityscapes + BUILDER_CONFIG_NAMES_TO_TEST = ['stereo_disparity_extra'] + SPLITS = { + 'train': 3, + 'train_extra': 4, + 'validation': 1, + } + + +if __name__ == '__main__': + testing.test_main() diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/disparity_trainextra.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/disparity_trainextra.zip new file mode 100644 index 00000000000..e5696d97c37 Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/disparity_trainextra.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/disparity_trainvaltest.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/disparity_trainvaltest.zip new file mode 100644 index 00000000000..98fd116e1af Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/disparity_trainvaltest.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/gtCoarse.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/gtCoarse.zip new file mode 100644 index 00000000000..0124e0198ea Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/gtCoarse.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/gtFine_trainvaltest.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/gtFine_trainvaltest.zip new file mode 100644 index 00000000000..2fd33b4493a Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/gtFine_trainvaltest.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/leftImg8bit_trainextra.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/leftImg8bit_trainextra.zip new file mode 100644 index 00000000000..7df50ad56df Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/leftImg8bit_trainextra.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/leftImg8bit_trainvaltest.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/leftImg8bit_trainvaltest.zip new file mode 100644 index 00000000000..a9d1b5ee4e3 Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/leftImg8bit_trainvaltest.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/rightImg8bit_trainextra.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/rightImg8bit_trainextra.zip new file mode 100644 index 00000000000..6b1d5f954a5 Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/rightImg8bit_trainextra.zip differ diff --git a/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/rightImg8bit_trainvaltest.zip b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/rightImg8bit_trainvaltest.zip new file mode 100644 index 00000000000..dd5472eec36 Binary files /dev/null and b/tensorflow_datasets/testing/test_data/fake_examples/cityscapes/rightImg8bit_trainvaltest.zip differ