fix expression

cvat-ai · Jun 8, 2020 · f21c824 · f21c824
1 parent f2b00ac
commit f21c824
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/datumaro/datumaro/plugins/voc_format/extractor.py b/datumaro/datumaro/plugins/voc_format/extractor.py
@@ -32,12 +32,15 @@ def __init__(self, path):
         super().__init__(subset=osp.splitext(osp.basename(path))[0])
 
         self._categories = self._load_categories(self._dataset_dir)
-        log.debug("Loaded labels: %s",
-            ', '.join("'%s' (%s, %s, %s)" % \
-                (l.name, *self._categories[AnnotationType.mask].colormap[idx])
-            for idx, l in enumerate(
-                self._categories[AnnotationType.label].items))
-        )
+
+        label_color = lambda label_idx: \
+            self._categories[AnnotationType.mask].colormap.get(label_idx, None)
+        log.debug("Loaded labels: %s" % ', '.join(
+            "'%s' %s" % (l.name, ('(%s, %s, %s)' % c) if c else '')
+            for i, l, c in ((i, l, label_color(i)) for i, l in enumerate(
+                self._categories[AnnotationType.label].items
+            ))
+        ))
         self._items = self._load_subset_list(path)
 
     def categories(self):