[Datumaro] Allow empty COCO dataset export (#1272)

* Allow empty dataset export in coco * Add line to changelog Co-authored-by: Nikita Manovich <[email protected]>
cvat-ai · Mar 18, 2020 · 9545132 · 9545132
1 parent 936d305
commit 9545132
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 -
 
 ### Fixed
+- Annotation-less tasks now can be exported as empty datasets in COCO ([#1277](https://github.com/opencv/cvat/issues/1277))
 - Frame name matching for video annotations import -
   allowed `frame_XXXXXX[.ext]` format ([#1274](https://github.com/opencv/cvat/pull/1274))
 

diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py
@@ -329,20 +329,24 @@ def save_categories(self, dataset):
         label_categories = dataset.categories().get(AnnotationType.label)
         if label_categories is None:
             return
-        points_categories = dataset.categories().get(AnnotationType.points)
-        if points_categories is None:
-            return
-
-        for idx, kp_cat in points_categories.items.items():
-            label_cat = label_categories.items[idx]
+        point_categories = dataset.categories().get(AnnotationType.points)
 
+        for idx, label_cat in enumerate(label_categories.items):
             cat = {
                 'id': 1 + idx,
                 'name': _cast(label_cat.name, str, ''),
                 'supercategory': _cast(label_cat.parent, str, ''),
-                'keypoints': [str(l) for l in kp_cat.labels],
-                'skeleton': [int(i) for i in kp_cat.adjacent],
+                'keypoints': [],
+                'skeleton': [],
             }
+
+            if point_categories is not None:
+                kp_cat = point_categories.items.get(idx)
+                if kp_cat is not None:
+                    cat.update({
+                        'keypoints': [str(l) for l in kp_cat.labels],
+                        'skeleton': [int(i) for i in kp_cat.adjacent],
+                    })
             self.categories.append(cat)
 
     def save_annotations(self, item):
@@ -447,14 +451,19 @@ class _Converter:
     def __init__(self, extractor, save_dir,
             tasks=None, save_images=False, segmentation_mode=None,
             crop_covered=False):
-        assert tasks is None or isinstance(tasks, (CocoTask, list))
+        assert tasks is None or isinstance(tasks, (CocoTask, list, str))
         if tasks is None:
             tasks = list(self._TASK_CONVERTER)
         elif isinstance(tasks, CocoTask):
             tasks = [tasks]
+        elif isinstance(tasks, str):
+            tasks = [CocoTask[tasks]]
         else:
-            for t in tasks:
-                assert t in CocoTask
+            for i, t in enumerate(tasks):
+                if isinstance(t, str):
+                    tasks[i] = CocoTask[t]
+                else:
+                    assert t in CocoTask, t
         self._tasks = tasks
 
         self._extractor = extractor
@@ -546,9 +555,8 @@ def convert(self):
                     task_conv.save_annotations(item)
 
             for task, task_conv in task_converters.items():
-                if not task_conv.is_empty():
-                    task_conv.write(osp.join(self._ann_dir,
-                        '%s_%s.json' % (task.name, subset_name)))
+                task_conv.write(osp.join(self._ann_dir,
+                    '%s_%s.json' % (task.name, subset_name)))
 
 class CocoConverter(Converter, CliPlugin):
     @staticmethod

diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py
@@ -632,10 +632,13 @@ def __iter__(self):
 
             def categories(self):
                 label_cat = LabelCategories()
+                point_cat = PointsCategories()
                 for label in range(10):
                     label_cat.add('label_' + str(label))
+                    point_cat.add(label)
                 return {
                     AnnotationType.label: label_cat,
+                    AnnotationType.points: point_cat,
                 }
 
         with TestDir() as test_dir:
@@ -651,4 +654,4 @@ def __iter__(self):
 
         with TestDir() as test_dir:
             self._test_save_and_load(TestExtractor(),
-                CocoConverter(), test_dir)
+                CocoConverter(tasks='image_info'), test_dir)