tools/accuracy_checker/accuracy_checker/annotation_converters/cityscapes.py

   1 from pathlib import Path
   2 from ..representation import SegmentationAnnotation
   3 from ..representation.segmentation_representation import GTMaskLoader
   4 from ..config import PathField, StringField, BoolField
   5 from .format_converter import BaseFormatConverter, BaseFormatConverterConfig
   6
   7
   8 train_meta = {
   9     'label_map': {
  10         0: 'road', 1: 'sidewalk', 2: 'building', 3: 'wall', 4: 'fence', 5: 'pole', 6: 'traffic light',
  11         7: 'traffic sign', 8: 'vegetation', 9: 'terrain', 10: 'sky', 11: 'person', 12: 'rider', 13: 'car',
  12         14: 'truck', 15: 'bus', 16: 'train', 17: 'motorcycle', 18: 'bicycle'
  13     },
  14     'segmentation_colors': (
  15         (128, 64, 128), (244, 35, 232), (70, 70, 70), (102, 102, 156), (190, 153, 153), (153, 153, 153),
  16         (250, 170, 30), (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0),
  17         (0, 0, 142), (0, 0, 70), (0, 60, 100), (0, 80, 100), (0, 0, 230), (119, 11, 32)
  18     ),
  19 }
  20
  21 full_dataset_meta = {
  22     'segmentation_colors' : (
  23         (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81), (128, 64, 128),
  24         (244, 35, 232), (250, 170, 160), (230, 150, 140), (70, 70, 70), (102, 102, 156), (190, 153, 153),
  25         (180, 165, 180), (150, 100, 100), (150, 120, 90), (153, 153, 153), (153, 153, 153), (250, 170, 30),
  26         (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0), (0, 0, 142),
  27         (0, 0, 70), (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32)
  28     ),
  29     'label_map': {
  30         0: 'unlabeled', 1:  'ego vehicle', 2: 'rectification border', 3: 'out of roi', 4: 'static', 5: 'dynamic',
  31         6: 'ground', 7: 'road', 8: 'sidewalk', 9: 'parking', 10: 'rail track', 11: 'building', 12: 'wall',
  32         13: 'fence', 14: 'guard rail', 15: 'bridge', 16: 'tunnel', 17: 'pole', 18: 'polegroup', 19: 'traffic light',
  33         20: 'traffic sign', 21: 'vegetation', 22: 'terrain', 23: 'sky', 24: 'person', 25: 'rider', 26: 'car',
  34         27: 'truck', 28: 'bus', 29: 'caravan', 30: 'trailer', 31: 'train', 32: 'motorcycle', 33: 'bicycle',
  35         -1: 'license plate'
  36     }
  37 }
  38
  39
  40 class CityscapesConverterConfig(BaseFormatConverterConfig):
  41     dataset_root_dir = PathField(is_directory=True)
  42     images_subfolder = StringField(optional=True)
  43     masks_subfolder = StringField(optional=True)
  44     masks_suffix = StringField(optional=True)
  45     images_suffix = StringField(optional=True)
  46     use_full_label_map = BoolField(optional=True)
  47
  48
  49 class CityscapesConverter(BaseFormatConverter):
  50     __provider__ = 'cityscapes'
  51
  52     _config_validator_type = CityscapesConverterConfig
  53
  54     def configure(self):
  55         self.dataset_root = self.config['dataset_root_dir']
  56         self.images_dir = self.config.get('images_subfolder', 'imgsFine/leftImg8bit/val')
  57         self.masks_dir = self.config.get('masks_subfolder', 'gtFine/val')
  58         self.masks_suffix = self.config.get('masks_suffix', '_gtFine_labelTrainIds')
  59         self.images_suffix = self.config.get('images_suffix', '_leftImg8bit')
  60         self.use_full_label_map = self.config.get('use_full_label_map', False)
  61
  62
  63     def convert(self):
  64         images = list(self.dataset_root.rglob(r'{}/*/*{}.png'.format(self.images_dir, self.images_suffix)))
  65         annotations = []
  66         for image in images:
  67             identifier = str(Path(self.images_dir).joinpath(*image.parts[-2:]))
  68             mask = Path(self.masks_dir) / image.parts[-2] / self.masks_suffix.join(
  69                 str(image.name).split(self.images_suffix)
  70             )
  71             annotations.append(SegmentationAnnotation(identifier, mask, mask_loader=GTMaskLoader.PILLOW))
  72
  73         return annotations, full_dataset_meta if self.use_full_label_map else train_meta