Fixed mypy issue and added tests and category infos

vfdev-5 · vfdev-5 · commit e003b714d8fb · 2022-02-21T14:33:32.000Z
diff --git 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
@@ -1,6 +1,7 @@
+from collections import namedtuple
 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,8 +11,9 @@
     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
+from torchvision.prototype.utils._internal import FrozenMapping
 
 
 class CityscapesDatasetInfo(DatasetInfo):
@@ -43,20 +45,66 @@ def __init__(self, **kwargs: Any) -> None:
         )
 
 
+CityscapesClass = namedtuple(
+    "CityscapesClass",
+    ["name", "id", "train_id", "category", "category_id", "has_instances", "ignore_in_eval", "color"],
+)
+
+
 class Cityscapes(Dataset):
+
+    categories_to_details: FrozenMapping = FrozenMapping(
+        {
+            "unlabeled": CityscapesClass("unlabeled", 0, 255, "void", 0, False, True, (0, 0, 0)),
+            "ego vehicle": CityscapesClass("ego vehicle", 1, 255, "void", 0, False, True, (0, 0, 0)),
+            "rectification border": CityscapesClass("rectification border", 2, 255, "void", 0, False, True, (0, 0, 0)),
+            "out of roi": CityscapesClass("out of roi", 3, 255, "void", 0, False, True, (0, 0, 0)),
+            "static": CityscapesClass("static", 4, 255, "void", 0, False, True, (0, 0, 0)),
+            "dynamic": CityscapesClass("dynamic", 5, 255, "void", 0, False, True, (111, 74, 0)),
+            "ground": CityscapesClass("ground", 6, 255, "void", 0, False, True, (81, 0, 81)),
+            "road": CityscapesClass("road", 7, 0, "flat", 1, False, False, (128, 64, 128)),
+            "sidewalk": CityscapesClass("sidewalk", 8, 1, "flat", 1, False, False, (244, 35, 232)),
+            "parking": CityscapesClass("parking", 9, 255, "flat", 1, False, True, (250, 170, 160)),
+            "rail track": CityscapesClass("rail track", 10, 255, "flat", 1, False, True, (230, 150, 140)),
+            "building": CityscapesClass("building", 11, 2, "construction", 2, False, False, (70, 70, 70)),
+            "wall": CityscapesClass("wall", 12, 3, "construction", 2, False, False, (102, 102, 156)),
+            "fence": CityscapesClass("fence", 13, 4, "construction", 2, False, False, (190, 153, 153)),
+            "guard rail": CityscapesClass("guard rail", 14, 255, "construction", 2, False, True, (180, 165, 180)),
+            "bridge": CityscapesClass("bridge", 15, 255, "construction", 2, False, True, (150, 100, 100)),
+            "tunnel": CityscapesClass("tunnel", 16, 255, "construction", 2, False, True, (150, 120, 90)),
+            "pole": CityscapesClass("pole", 17, 5, "object", 3, False, False, (153, 153, 153)),
+            "polegroup": CityscapesClass("polegroup", 18, 255, "object", 3, False, True, (153, 153, 153)),
+            "traffic light": CityscapesClass("traffic light", 19, 6, "object", 3, False, False, (250, 170, 30)),
+            "traffic sign": CityscapesClass("traffic sign", 20, 7, "object", 3, False, False, (220, 220, 0)),
+            "vegetation": CityscapesClass("vegetation", 21, 8, "nature", 4, False, False, (107, 142, 35)),
+            "terrain": CityscapesClass("terrain", 22, 9, "nature", 4, False, False, (152, 251, 152)),
+            "sky": CityscapesClass("sky", 23, 10, "sky", 5, False, False, (70, 130, 180)),
+            "person": CityscapesClass("person", 24, 11, "human", 6, True, False, (220, 20, 60)),
+            "rider": CityscapesClass("rider", 25, 12, "human", 6, True, False, (255, 0, 0)),
+            "car": CityscapesClass("car", 26, 13, "vehicle", 7, True, False, (0, 0, 142)),
+            "truck": CityscapesClass("truck", 27, 14, "vehicle", 7, True, False, (0, 0, 70)),
+            "bus": CityscapesClass("bus", 28, 15, "vehicle", 7, True, False, (0, 60, 100)),
+            "caravan": CityscapesClass("caravan", 29, 255, "vehicle", 7, True, True, (0, 0, 90)),
+            "trailer": CityscapesClass("trailer", 30, 255, "vehicle", 7, True, True, (0, 0, 110)),
+            "train": CityscapesClass("train", 31, 16, "vehicle", 7, True, False, (0, 80, 100)),
+            "motorcycle": CityscapesClass("motorcycle", 32, 17, "vehicle", 7, True, False, (0, 0, 230)),
+            "bicycle": CityscapesClass("bicycle", 33, 18, "vehicle", 7, True, False, (119, 11, 32)),
+            "license plate": CityscapesClass("license plate", -1, -1, "vehicle", 7, False, True, (0, 0, 142)),
+        }
+    )
+
     def _make_info(self) -> DatasetInfo:
         name = "cityscapes"
-        categories = None
 
         return CityscapesDatasetInfo(
             name,
-            categories=categories,
+            categories=list(self.categories_to_details.keys()),
             homepage="http://www.cityscapes-dataset.com/",
             valid_options=dict(
                 split=("train", "val", "test", "train_extra"),
                 mode=("fine", "coarse"),
-                # target_type=("instance", "semantic", "polygon", "color")
             ),
+            extra=dict(classname_to_details=self.categories_to_details),
         )
 
     _FILES_CHECKSUMS = {
@@ -67,8 +115,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 +127,22 @@ 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 +154,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 +163,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 +174,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,12 +199,12 @@ 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")]
             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]
@@ -179,4 +230,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)
