Use default colors for masks

cvat-ai · May 16, 2020 · 7e0a6df · 7e0a6df
1 parent ccd2fbd
commit 7e0a6df
Showing 3 changed files with 777 additions and 1 deletion.
diff --git a/cvat/apps/dataset_manager/formats/mask.py b/cvat/apps/dataset_manager/formats/mask.py
@@ -2,14 +2,17 @@
 #
 # SPDX-License-Identifier: MIT
 
+import os.path as osp
 from tempfile import TemporaryDirectory
 
 from pyunpack import Archive
 
 from cvat.apps.dataset_manager.bindings import (CvatTaskDataExtractor,
     import_dm_annotations)
 from cvat.apps.dataset_manager.util import make_zip_archive
+from datumaro.cli.util import make_file_name
 from datumaro.components.project import Dataset
+from datumaro.util.mask_tools import generate_colormap
 
 from .registry import dm_env, exporter, importer
 
@@ -25,7 +28,8 @@ def _export(dst_file, task_data, save_images=False):
     extractor = Dataset.from_extractors(extractor) # apply lazy transforms
     with TemporaryDirectory() as temp_dir:
         converter = dm_env.make_converter('voc_segmentation',
-            apply_colormap=True, label_map='source', save_images=save_images)
+            apply_colormap=True, label_map=make_colormap(task_data),
+            save_images=save_images)
         converter(extractor, save_dir=temp_dir)
 
         make_zip_archive(temp_dir, dst_file)
@@ -39,3 +43,38 @@ def _import(src_file, task_data):
         masks_to_polygons = dm_env.transforms.get('masks_to_polygons')
         dataset = dataset.transform(masks_to_polygons)
         import_dm_annotations(dataset, task_data)
+
+
+DEFAULT_COLORMAP_CAPACITY = 2000
+DEFAULT_COLORMAP_PATH = osp.join(osp.dirname(__file__), 'default_colors.txt')
+def parse_default_colors(file_path=None):
+    file_path = file_path or DEFAULT_COLORMAP_PATH
+
+    colors = {}
+    with open(file_path) as f:
+        for line in f:
+            line = line.strip()
+            if not line or line[0] == '#':
+                continue
+            _, label, color = line.split(':')
+            colors[label] = tuple(map(int, color.split(',')))
+    return colors
+
+def normalize_label(label):
+    label = make_file_name(label) # basically, convert to ASCII lowercase
+    label = label.replace('-', '_')
+    return label
+
+def make_colormap(task_data):
+    labels = sorted(get_labels(task_data))
+    predefined = parse_default_colors()
+
+    # NOTE: using pop() to avoid collisions
+    colormap = {k: predefined.pop(normalize_label(k), None) for k in labels}
+
+    random_labels = [k for k in labels if not colormap[k]]
+    if random_labels:
+        colors = generate_colormap(DEFAULT_COLORMAP_CAPACITY + len(random_labels))
+        for i, label in enumerate(random_labels):
+            colormap[label] = colors[DEFAULT_COLORMAP_CAPACITY + i]
+    return {l: [c, [], []] for l, c in colormap.items()}