From 3604f0c5ce8b3acc07f048817d2f2373544290f9 Mon Sep 17 00:00:00 2001 From: zhiltsov-max Date: Fri, 28 Feb 2020 22:30:20 +0300 Subject: [PATCH] Auto-generate labelmap for voc from task (#1214) --- cvat/apps/annotation/pascal_voc.py | 8 ++++++++ datumaro/datumaro/plugins/voc_format/extractor.py | 7 ++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/cvat/apps/annotation/pascal_voc.py b/cvat/apps/annotation/pascal_voc.py index e65a186b..801d1ba4 100644 --- a/cvat/apps/annotation/pascal_voc.py +++ b/cvat/apps/annotation/pascal_voc.py @@ -36,6 +36,14 @@ def load(file_object, annotations): with TemporaryDirectory() as tmp_dir: Archive(archive_file).extractall(tmp_dir) + # put label map from the task if not present + labelmap_file = osp.join(tmp_dir, 'labelmap.txt') + if not osp.isfile(labelmap_file): + labels = (label['name'] + ':::' + for _, label in annotations.meta['task']['labels']) + with open(labelmap_file, 'w') as f: + f.write('\n'.join(labels)) + # support flat archive layout anno_dir = osp.join(tmp_dir, 'Annotations') if not osp.isdir(anno_dir): diff --git a/datumaro/datumaro/plugins/voc_format/extractor.py b/datumaro/datumaro/plugins/voc_format/extractor.py index 87a3374f..31047e85 100644 --- a/datumaro/datumaro/plugins/voc_format/extractor.py +++ b/datumaro/datumaro/plugins/voc_format/extractor.py @@ -155,7 +155,12 @@ class VocExtractor(SourceExtractor): def _get_label_id(self, label): label_id, _ = self._categories[AnnotationType.label].find(label) - assert label_id is not None + if label_id is None: + log.debug("Unknown label '%s'. Loaded labels: %s", + label, + ', '.join("'%s'" % s.name + for s in self._categories[AnnotationType.label].items)) + raise Exception("Unknown label '%s'" % label) return label_id @staticmethod