diff --git a/test/builtin_dataset_mocks.py b/test/builtin_dataset_mocks.py index 123d8f29d3f..9f4f8a54167 100644 --- a/test/builtin_dataset_mocks.py +++ b/test/builtin_dataset_mocks.py @@ -1344,3 +1344,96 @@ def pcam(info, root, config): compressed_file.write(compressed_data) return num_images + + +class CityScapesMockData: + + _ARCHIVE_NAMES = { + ("Coarse", "train"): [("gtCoarse.zip", "gtCoarse"), ("leftImg8bit_trainvaltest.zip", "leftImg8bit")], + ("Coarse", "train_extra"): [("gtCoarse.zip", "gtCoarse"), ("leftImg8bit_trainextra.zip", "leftImg8bit")], + ("Coarse", "val"): [("gtCoarse.zip", "gtCoarse"), ("leftImg8bit_trainvaltest.zip", "leftImg8bit")], + ("Fine", "train"): [("gtFine_trainvaltest.zip", "gtFine"), ("leftImg8bit_trainvaltest.zip", "leftImg8bit")], + ("Fine", "test"): [("gtFine_trainvaltest.zip", "gtFine"), ("leftImg8bit_trainvaltest.zip", "leftImg8bit")], + ("Fine", "val"): [("gtFine_trainvaltest.zip", "gtFine"), ("leftImg8bit_trainvaltest.zip", "leftImg8bit")], + } + + @classmethod + def generate(cls, root, config): + + mode = config.mode.capitalize() + split = config.split + + if split in ["train", "train_extra"]: + cities = ["bochum", "bremen"] + num_samples = 3 + else: + cities = ["bochum"] + num_samples = 2 + + polygon_target = { + "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], + ], + }, + ], + } + + gt_dir = root / f"gt{mode}" + + for city in cities: + + def make_image(name, size=10): + create_image_folder( + root=gt_dir / split, + name=city, + file_name_fn=lambda idx: name.format(idx=idx), + size=size, + num_examples=num_samples, + ) + + make_image(f"{city}_000000_00000" + "{idx}" + f"_gt{mode}_instanceIds.png") + make_image(f"{city}_000000_00000" + "{idx}" + f"_gt{mode}_labelIds.png") + make_image(f"{city}_000000_00000" + "{idx}" + f"_gt{mode}_color.png", size=(4, 10, 10)) + + for idx in range(num_samples): + polygon_target_name = gt_dir / split / city / f"{city}_000000_00000{idx}_gt{mode}_polygons.json" + with open(polygon_target_name, "w") as outfile: + json.dump(polygon_target, outfile) + + # Create leftImg8bit folder + for city in cities: + create_image_folder( + root=root / "leftImg8bit" / split, + name=city, + file_name_fn=lambda idx: f"{city}_000000_00000{idx}_leftImg8bit.png", + num_examples=num_samples, + ) + + for zip_name, folder_name in cls._ARCHIVE_NAMES[(mode, split)]: + make_zip(root, zip_name, folder_name) + return len(cities) * num_samples + + +@register_mock +def cityscapes(info, root, config): + return CityScapesMockData.generate(root, config) diff --git a/torchvision/prototype/datasets/_builtin/cityscapes.py b/torchvision/prototype/datasets/_builtin/cityscapes.py index 0fbae6ec464..6d9128aeaff 100644 --- a/torchvision/prototype/datasets/_builtin/cityscapes.py +++ b/torchvision/prototype/datasets/_builtin/cityscapes.py @@ -1,6 +1,6 @@ from functools import partial from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional, Tuple from torchdata.datapipes.iter import IterDataPipe, Mapper, Filter, Demultiplexer, IterKeyZipper, JsonParser from torchvision.prototype.datasets.utils import ( @@ -10,7 +10,7 @@ ManualDownloadResource, OnlineResource, ) -from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE +from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE, hint_sharding, hint_shuffling from torchvision.prototype.features import EncodedImage @@ -55,7 +55,6 @@ def _make_info(self) -> DatasetInfo: valid_options=dict( split=("train", "val", "test", "train_extra"), mode=("fine", "coarse"), - # target_type=("instance", "semantic", "polygon", "color") ), ) @@ -67,8 +66,9 @@ def _make_info(self) -> DatasetInfo: } def resources(self, config: DatasetConfig) -> List[OnlineResource]: + resources: List[OnlineResource] = [] if config.mode == "fine": - resources = [ + resources += [ CityscapesResource( file_name="leftImg8bit_trainvaltest.zip", sha256=self._FILES_CHECKSUMS["leftImg8bit_trainvaltest.zip"], @@ -78,20 +78,21 @@ def resources(self, config: DatasetConfig) -> List[OnlineResource]: ), ] else: - resources = [ + split_label = "trainextra" if config.split == "train_extra" else "trainvaltest" + resources += [ CityscapesResource( - file_name="leftImg8bit_trainextra.zip", sha256=self._FILES_CHECKSUMS["leftImg8bit_trainextra.zip"] + file_name=f"leftImg8bit_{split_label}.zip", sha256=self._FILES_CHECKSUMS[f"leftImg8bit_{split_label}.zip"] ), CityscapesResource(file_name="gtCoarse.zip", sha256=self._FILES_CHECKSUMS["gtCoarse.zip"]), ] return resources - def _filter_split_images(self, data, *, req_split: str): + def _filter_split_images(self, data: Tuple[str, Any], *, req_split: str) -> bool: path = Path(data[0]) split = path.parent.parts[-2] return split == req_split and ".png" == path.suffix - def _filter_classify_targets(self, data, *, req_split: str): + def _filter_classify_targets(self, data: Tuple[str, Any], *, req_split: str) -> Optional[int]: path = Path(data[0]) name = path.name split = path.parent.parts[-2] @@ -103,7 +104,7 @@ def _filter_classify_targets(self, data, *, req_split: str): return i return None - def _prepare_sample(self, data): + def _prepare_sample(self, data: Tuple[Tuple[str, Any], Any]) -> Dict[str, Any]: (img_path, img_data), target_data = data color_path, color_data = target_data[1] @@ -112,7 +113,7 @@ def _prepare_sample(self, data): target_data = target_data[0] label_path, label_data = target_data[1] target_data = target_data[0] - instance_path, instance_data = target_data + instances_path, instance_data = target_data return dict( image_path=img_path, @@ -123,7 +124,7 @@ def _prepare_sample(self, data): polygon=polygon_data, segmentation_path=label_path, segmentation=EncodedImage.from_file(label_data), - instances_path=color_path, + instances_path=instances_path, instances=EncodedImage.from_file(instance_data), ) @@ -148,18 +149,20 @@ def _make_datapipe( # targets_dps[2] is for json polygon, we have to decode them targets_dps[2] = JsonParser(targets_dps[2]) - def img_key_fn(data): + def img_key_fn(data: Tuple[str, Any]) -> str: stem = Path(data[0]).stem stem = stem[: -len("_leftImg8bit")] + print("img_key stem:", stem, "<-", Path(data[0]).name) return stem - def target_key_fn(data, level=0): + def target_key_fn(data: Tuple[Any, Any], level: int = 0) -> str: path = data[0] for _ in range(level): path = path[0] stem = Path(path).stem i = stem.rfind("_gt") stem = stem[:i] + print("target_key stem:", stem, level, "<-", Path(path).name) return stem zipped_targets_dp = targets_dps[0] @@ -179,4 +182,6 @@ def target_key_fn(data, level=0): ref_key_fn=partial(target_key_fn, level=len(targets_dps) - 1), buffer_size=INFINITE_BUFFER_SIZE, ) + samples = hint_sharding(samples) + samples = hint_shuffling(samples) return Mapper(samples, fn=self._prepare_sample)