|
| 1 | + |
| 2 | +# Copyright (C) 2020 Intel Corporation |
| 3 | +# |
| 4 | +# SPDX-License-Identifier: MIT |
| 5 | + |
| 6 | +import os |
| 7 | +import os.path as osp |
| 8 | +from collections import OrderedDict |
| 9 | +from enum import Enum |
| 10 | +from glob import glob |
| 11 | + |
| 12 | +import numpy as np |
| 13 | +from datumaro.components.converter import Converter |
| 14 | +from datumaro.components.extractor import (AnnotationType, CompiledMask, |
| 15 | + DatasetItem, Importer, LabelCategories, Mask, |
| 16 | + MaskCategories, SourceExtractor) |
| 17 | +from datumaro.util import find, str_to_bool |
| 18 | +from datumaro.util.image import save_image |
| 19 | +from datumaro.util.mask_tools import lazy_mask, paint_mask, generate_colormap |
| 20 | + |
| 21 | + |
| 22 | +CamvidLabelMap = OrderedDict([ |
| 23 | + ('Void', (0, 0, 0)), |
| 24 | + ('Animal', (64, 128, 64)), |
| 25 | + ('Archway', (192, 0, 128)), |
| 26 | + ('Bicyclist', (0, 128, 192)), |
| 27 | + ('Bridge', (0, 128, 64)), |
| 28 | + ('Building', (128, 0, 0)), |
| 29 | + ('Car', (64, 0, 128)), |
| 30 | + ('CartLuggagePram', (64, 0, 192)), |
| 31 | + ('Child', (192, 128, 64)), |
| 32 | + ('Column_Pole', (192, 192, 128)), |
| 33 | + ('Fence', (64, 64, 128)), |
| 34 | + ('LaneMkgsDriv', (128, 0, 192)), |
| 35 | + ('LaneMkgsNonDriv', (192, 0, 64)), |
| 36 | + ('Misc_Text', (128, 128, 64)), |
| 37 | + ('MotorcycycleScooter', (192, 0, 192)), |
| 38 | + ('OtherMoving', (128, 64, 64)), |
| 39 | + ('ParkingBlock', (64, 192, 128)), |
| 40 | + ('Pedestrian', (64, 64, 0)), |
| 41 | + ('Road', (128, 64, 128)), |
| 42 | + ('RoadShoulder', (128, 128, 192)), |
| 43 | + ('Sidewalk', (0, 0, 192)), |
| 44 | + ('SignSymbol', (192, 128, 128)), |
| 45 | + ('Sky', (128, 128, 128)), |
| 46 | + ('SUVPickupTruck', (64, 128, 192)), |
| 47 | + ('TrafficCone', (0, 0, 64)), |
| 48 | + ('TrafficLight', (0, 64, 64)), |
| 49 | + ('Train', (192, 64, 128)), |
| 50 | + ('Tree', (128, 128, 0)), |
| 51 | + ('Truck_Bus', (192, 128, 192)), |
| 52 | + ('Tunnel', (64, 0, 64)), |
| 53 | + ('VegetationMisc', (192, 192, 0)), |
| 54 | + ('Wall', (64, 192, 0)) |
| 55 | +]) |
| 56 | + |
| 57 | +class CamvidPath: |
| 58 | + LABELMAP_FILE = 'label_colors.txt' |
| 59 | + SEGM_DIR = "annot" |
| 60 | + IMAGE_EXT = '.png' |
| 61 | + |
| 62 | + |
| 63 | +def parse_label_map(path): |
| 64 | + if not path: |
| 65 | + return None |
| 66 | + |
| 67 | + label_map = OrderedDict() |
| 68 | + with open(path, 'r') as f: |
| 69 | + for line in f: |
| 70 | + # skip empty and commented lines |
| 71 | + line = line.strip() |
| 72 | + if not line or line and line[0] == '#': |
| 73 | + continue |
| 74 | + |
| 75 | + # color, name |
| 76 | + label_desc = line.strip().split() |
| 77 | + |
| 78 | + if 2 < len(label_desc): |
| 79 | + name = label_desc[3] |
| 80 | + color = tuple([int(c) for c in label_desc[:-1]]) |
| 81 | + else: |
| 82 | + name = label_desc[0] |
| 83 | + color = None |
| 84 | + |
| 85 | + if name in label_map: |
| 86 | + raise ValueError("Label '%s' is already defined" % name) |
| 87 | + |
| 88 | + label_map[name] = color |
| 89 | + return label_map |
| 90 | + |
| 91 | +def write_label_map(path, label_map): |
| 92 | + with open(path, 'w') as f: |
| 93 | + for label_name, label_desc in label_map.items(): |
| 94 | + if label_desc: |
| 95 | + color_rgb = ' '.join(str(c) for c in label_desc) |
| 96 | + else: |
| 97 | + color_rgb = '' |
| 98 | + f.write('%s %s\n' % (color_rgb, label_name)) |
| 99 | + |
| 100 | +def make_camvid_categories(label_map=None): |
| 101 | + if label_map is None: |
| 102 | + label_map = CamvidLabelMap |
| 103 | + |
| 104 | + # There must always be a label with color (0, 0, 0) at index 0 |
| 105 | + bg_label = find(label_map.items(), lambda x: x[1] == (0, 0, 0)) |
| 106 | + if bg_label is not None: |
| 107 | + bg_label = bg_label[0] |
| 108 | + else: |
| 109 | + bg_label = 'background' |
| 110 | + if bg_label not in label_map: |
| 111 | + has_colors = any(v is not None for v in label_map.values()) |
| 112 | + color = (0, 0, 0) if has_colors else None |
| 113 | + label_map[bg_label] = color |
| 114 | + label_map.move_to_end(bg_label, last=False) |
| 115 | + |
| 116 | + categories = {} |
| 117 | + label_categories = LabelCategories() |
| 118 | + for label, desc in label_map.items(): |
| 119 | + label_categories.add(label) |
| 120 | + categories[AnnotationType.label] = label_categories |
| 121 | + |
| 122 | + has_colors = any(v is not None for v in label_map.values()) |
| 123 | + if not has_colors: # generate new colors |
| 124 | + colormap = generate_colormap(len(label_map)) |
| 125 | + else: # only copy defined colors |
| 126 | + label_id = lambda label: label_categories.find(label)[0] |
| 127 | + colormap = { label_id(name): (desc[0], desc[1], desc[2]) |
| 128 | + for name, desc in label_map.items() } |
| 129 | + mask_categories = MaskCategories(colormap) |
| 130 | + mask_categories.inverse_colormap # pylint: disable=pointless-statement |
| 131 | + categories[AnnotationType.mask] = mask_categories |
| 132 | + return categories |
| 133 | + |
| 134 | + |
| 135 | +class CamvidExtractor(SourceExtractor): |
| 136 | + def __init__(self, path): |
| 137 | + assert osp.isfile(path), path |
| 138 | + self._path = path |
| 139 | + self._dataset_dir = osp.dirname(path) |
| 140 | + super().__init__(subset=osp.splitext(osp.basename(path))[0]) |
| 141 | + |
| 142 | + self._categories = self._load_categories(self._dataset_dir) |
| 143 | + self._items = list(self._load_items(path).values()) |
| 144 | + |
| 145 | + def _load_categories(self, path): |
| 146 | + label_map = None |
| 147 | + label_map_path = osp.join(path, CamvidPath.LABELMAP_FILE) |
| 148 | + if osp.isfile(label_map_path): |
| 149 | + label_map = parse_label_map(label_map_path) |
| 150 | + else: |
| 151 | + label_map = CamvidLabelMap |
| 152 | + self._labels = [label for label in label_map] |
| 153 | + return make_camvid_categories(label_map) |
| 154 | + |
| 155 | + def _load_items(self, path): |
| 156 | + items = {} |
| 157 | + with open(path, encoding='utf-8') as f: |
| 158 | + for line in f: |
| 159 | + objects = line.split() |
| 160 | + image = objects[0] |
| 161 | + item_id = ('/'.join(image.split('/')[2:]))[:-len(CamvidPath.IMAGE_EXT)] |
| 162 | + image_path = osp.join(self._dataset_dir, |
| 163 | + (image, image[1:])[image[0] == '/']) |
| 164 | + item_annotations = [] |
| 165 | + if 1 < len(objects): |
| 166 | + gt = objects[1] |
| 167 | + gt_path = osp.join(self._dataset_dir, |
| 168 | + (gt, gt[1:]) [gt[0] == '/']) |
| 169 | + inverse_cls_colormap = \ |
| 170 | + self._categories[AnnotationType.mask].inverse_colormap |
| 171 | + mask = lazy_mask(gt_path, inverse_cls_colormap) |
| 172 | + # loading mask through cache |
| 173 | + mask = mask() |
| 174 | + classes = np.unique(mask) |
| 175 | + labels = self._categories[AnnotationType.label]._indices |
| 176 | + labels = { labels[label_name]: label_name |
| 177 | + for label_name in labels } |
| 178 | + for label_id in classes: |
| 179 | + if labels[label_id] in self._labels: |
| 180 | + image = self._lazy_extract_mask(mask, label_id) |
| 181 | + item_annotations.append(Mask(image=image, label=label_id)) |
| 182 | + items[item_id] = DatasetItem(id=item_id, subset=self._subset, |
| 183 | + image=image_path, annotations=item_annotations) |
| 184 | + return items |
| 185 | + |
| 186 | + @staticmethod |
| 187 | + def _lazy_extract_mask(mask, c): |
| 188 | + return lambda: mask == c |
| 189 | + |
| 190 | + |
| 191 | +class CamvidImporter(Importer): |
| 192 | + @classmethod |
| 193 | + def find_sources(cls, path): |
| 194 | + subset_paths = [p for p in glob(osp.join(path, '**.txt'), recursive=True) |
| 195 | + if osp.basename(p) != CamvidPath.LABELMAP_FILE] |
| 196 | + sources = [] |
| 197 | + for subset_path in subset_paths: |
| 198 | + sources += cls._find_sources_recursive( |
| 199 | + subset_path, '.txt', 'camvid') |
| 200 | + return sources |
| 201 | + |
| 202 | + |
| 203 | +LabelmapType = Enum('LabelmapType', ['camvid', 'source']) |
| 204 | + |
| 205 | +class CamvidConverter(Converter): |
| 206 | + DEFAULT_IMAGE_EXT = '.png' |
| 207 | + |
| 208 | + @classmethod |
| 209 | + def build_cmdline_parser(cls, **kwargs): |
| 210 | + parser = super().build_cmdline_parser(**kwargs) |
| 211 | + |
| 212 | + parser.add_argument('--apply-colormap', type=str_to_bool, default=True, |
| 213 | + help="Use colormap for class masks (default: %(default)s)") |
| 214 | + parser.add_argument('--label-map', type=cls._get_labelmap, default=None, |
| 215 | + help="Labelmap file path or one of %s" % \ |
| 216 | + ', '.join(t.name for t in LabelmapType)) |
| 217 | + |
| 218 | + def __init__(self, extractor, save_dir, |
| 219 | + apply_colormap=True, label_map=None, **kwargs): |
| 220 | + super().__init__(extractor, save_dir, **kwargs) |
| 221 | + |
| 222 | + self._apply_colormap = apply_colormap |
| 223 | + |
| 224 | + if label_map is None: |
| 225 | + label_map = LabelmapType.source.name |
| 226 | + self._load_categories(label_map) |
| 227 | + |
| 228 | + def apply(self): |
| 229 | + subset_dir = self._save_dir |
| 230 | + os.makedirs(subset_dir, exist_ok=True) |
| 231 | + |
| 232 | + for subset_name, subset in self._extractor.subsets().items(): |
| 233 | + segm_list = {} |
| 234 | + for item in subset: |
| 235 | + masks = [a for a in item.annotations |
| 236 | + if a.type == AnnotationType.mask] |
| 237 | + |
| 238 | + if masks: |
| 239 | + compiled_mask = CompiledMask.from_instance_masks(masks, |
| 240 | + instance_labels=[self._label_id_mapping(m.label) |
| 241 | + for m in masks]) |
| 242 | + |
| 243 | + self.save_segm(osp.join(subset_dir, |
| 244 | + subset_name + CamvidPath.SEGM_DIR, |
| 245 | + item.id + CamvidPath.IMAGE_EXT), |
| 246 | + compiled_mask.class_mask) |
| 247 | + segm_list[item.id] = True |
| 248 | + else: |
| 249 | + segm_list[item.id] = False |
| 250 | + |
| 251 | + if self._save_images: |
| 252 | + self._save_image(item, osp.join(subset_dir, subset_name, |
| 253 | + item.id + CamvidPath.IMAGE_EXT)) |
| 254 | + |
| 255 | + self.save_segm_lists(subset_name, segm_list) |
| 256 | + self.save_label_map() |
| 257 | + |
| 258 | + def save_segm(self, path, mask, colormap=None): |
| 259 | + if self._apply_colormap: |
| 260 | + if colormap is None: |
| 261 | + colormap = self._categories[AnnotationType.mask].colormap |
| 262 | + mask = paint_mask(mask, colormap) |
| 263 | + save_image(path, mask, create_dir=True) |
| 264 | + |
| 265 | + def save_segm_lists(self, subset_name, segm_list): |
| 266 | + if not segm_list: |
| 267 | + return |
| 268 | + |
| 269 | + ann_file = osp.join(self._save_dir, subset_name + '.txt') |
| 270 | + with open(ann_file, 'w') as f: |
| 271 | + for item in segm_list: |
| 272 | + if segm_list[item]: |
| 273 | + path_mask = '/%s/%s' % (subset_name + CamvidPath.SEGM_DIR, |
| 274 | + item + CamvidPath.IMAGE_EXT) |
| 275 | + else: |
| 276 | + path_mask = '' |
| 277 | + f.write('/%s/%s %s\n' % (subset_name, |
| 278 | + item + CamvidPath.IMAGE_EXT, path_mask)) |
| 279 | + |
| 280 | + def save_label_map(self): |
| 281 | + path = osp.join(self._save_dir, CamvidPath.LABELMAP_FILE) |
| 282 | + labels = self._extractor.categories()[AnnotationType.label]._indices |
| 283 | + if len(self._label_map) > len(labels): |
| 284 | + self._label_map.pop('background') |
| 285 | + write_label_map(path, self._label_map) |
| 286 | + |
| 287 | + def _load_categories(self, label_map_source): |
| 288 | + if label_map_source == LabelmapType.camvid.name: |
| 289 | + # use the default Camvid colormap |
| 290 | + label_map = CamvidLabelMap |
| 291 | + |
| 292 | + elif label_map_source == LabelmapType.source.name and \ |
| 293 | + AnnotationType.mask not in self._extractor.categories(): |
| 294 | + # generate colormap for input labels |
| 295 | + labels = self._extractor.categories() \ |
| 296 | + .get(AnnotationType.label, LabelCategories()) |
| 297 | + label_map = OrderedDict((item.name, None) |
| 298 | + for item in labels.items) |
| 299 | + |
| 300 | + elif label_map_source == LabelmapType.source.name and \ |
| 301 | + AnnotationType.mask in self._extractor.categories(): |
| 302 | + # use source colormap |
| 303 | + labels = self._extractor.categories()[AnnotationType.label] |
| 304 | + colors = self._extractor.categories()[AnnotationType.mask] |
| 305 | + label_map = OrderedDict() |
| 306 | + for idx, item in enumerate(labels.items): |
| 307 | + color = colors.colormap.get(idx) |
| 308 | + if color is not None: |
| 309 | + label_map[item.name] = color |
| 310 | + |
| 311 | + elif isinstance(label_map_source, dict): |
| 312 | + label_map = OrderedDict( |
| 313 | + sorted(label_map_source.items(), key=lambda e: e[0])) |
| 314 | + |
| 315 | + elif isinstance(label_map_source, str) and osp.isfile(label_map_source): |
| 316 | + label_map = parse_label_map(label_map_source) |
| 317 | + |
| 318 | + else: |
| 319 | + raise Exception("Wrong labelmap specified, " |
| 320 | + "expected one of %s or a file path" % \ |
| 321 | + ', '.join(t.name for t in LabelmapType)) |
| 322 | + |
| 323 | + self._categories = make_camvid_categories(label_map) |
| 324 | + self._label_map = label_map |
| 325 | + self._label_id_mapping = self._make_label_id_map() |
| 326 | + |
| 327 | + def _make_label_id_map(self): |
| 328 | + source_labels = { |
| 329 | + id: label.name for id, label in |
| 330 | + enumerate(self._extractor.categories().get( |
| 331 | + AnnotationType.label, LabelCategories()).items) |
| 332 | + } |
| 333 | + target_labels = { |
| 334 | + label.name: id for id, label in |
| 335 | + enumerate(self._categories[AnnotationType.label].items) |
| 336 | + } |
| 337 | + id_mapping = { |
| 338 | + src_id: target_labels.get(src_label, 0) |
| 339 | + for src_id, src_label in source_labels.items() |
| 340 | + } |
| 341 | + |
| 342 | + def map_id(src_id): |
| 343 | + return id_mapping.get(src_id, 0) |
| 344 | + return map_id |
0 commit comments