add demo

datasets/cityscapes.py

@@ -0,0 +1,106 @@
+import torchvision
+import numpy as np
+from PIL import Image
+from typing import List, Any, Callable, Tuple
+from collections import namedtuple
+
+def get_cs_labeldata():
+    cls_names = ['road', 'sidewalk', 'parking', 'rail track', 'building',
+           'wall', 'fence', 'guard rail', 'bridge', 'tunnel',
+           'pole', 'polegroup', 'traffic light', 'traffic sign', 'vegetation',
+           'terrain', 'sky', 'person', 'rider', 'car',
+           'truck', 'bus', 'caravan', 'trailer', 'train',
+           'motorcycle', 'bicycle']
+    colormap = np.array([
+            [128, 64, 128],
+            [244, 35, 232],
+            [250, 170, 160],
+            [230, 150, 140],
+            [70, 70, 70],
+            [102, 102, 156],
+            [190, 153, 153],
+            [180, 165, 180],
+            [150, 100, 100],
+            [150, 120, 90],
+            [153, 153, 153],
+            [153, 153, 153],
+            [250, 170, 30],
+            [220, 220, 0],
+            [107, 142, 35],
+            [152, 251, 152],
+            [70, 130, 180],
+            [220, 20, 60],
+            [255, 0, 0],
+            [0, 0, 142],
+            [0, 0, 70],
+            [0, 60, 100],
+            [0, 0, 90],
+            [0, 0, 110],
+            [0, 80, 100],
+            [0, 0, 230],
+            [119, 11, 32],
+            [0, 0, 0],
+            [220, 220, 220]])
+    return cls_names, colormap
+
+class CityscapesDataset(torchvision.datasets.Cityscapes):
+
+    def __init__(self,
+                 transforms: List[Callable],
+                 *args: Any,
+                 **kwargs: Any):
+
+        super(CityscapesDataset, self).__init__(*args,
+                                                **kwargs,
+                                                target_type="semantic")
+        self.transforms = transforms
+        self.classes = ['road', 'sidewalk', 'parking', 'rail track', 'building',
+           'wall', 'fence', 'guard rail', 'bridge', 'tunnel',
+           'pole', 'polegroup', 'traffic light', 'traffic sign', 'vegetation',
+           'terrain', 'sky', 'person', 'rider', 'car',
+           'truck', 'bus', 'caravan', 'trailer', 'train',
+           'motorcycle', 'bicycle']
+        
+    def __getitem__(self, index: int) -> Tuple[Any, Any]:
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is a tuple of all target types if target_type is a list with more
+            than one item. Otherwise target is a json object if target_type="polygon", else the image segmentation.
+        """
+        img_pth = self.images[index]
+        image = Image.open(self.images[index]).convert('RGB')
+
+        targets: Any = []
+        for i, t in enumerate(self.target_type):
+            if t == 'polygon':
+                target = self._load_json(self.targets[index][i])
+            else:
+                target = Image.open(self.targets[index][i])
+
+            targets.append(target)
+
+        target = tuple(targets) if len(targets) > 1 else targets[0]
+
+        if self.transforms is not None:
+            image, target = self.transforms(image, target)
+
+        return image, target, img_pth
+
+def cityscapes(root: str,
+               split: str,
+               transforms: List[Callable]):
+    return CityscapesDataset(root=root,
+                             split=split,
+                             transforms=transforms)
+
+CityscapesClass = namedtuple('CityscapesClass', ['name', 'id', 'train_id', 'category', 'category_id',
+                                                    'has_instances', 'ignore_in_eval', 'color'])
+
+classes = ['road', 'sidewalk', 'parking', 'rail track', 'building',
+           'wall', 'fence', 'guard rail', 'bridge', 'tunnel',
+           'pole', 'polegroup', 'traffic light', 'traffic sign', 'vegetation',
+           'terrain', 'sky', 'person', 'rider', 'car',
+           'truck', 'bus', 'caravan', 'trailer', 'train',
+           'motorcycle', 'bicycle']