COCO: when saving, don't drop standalone skeletons that have group IDs

Previously, the `find_instance_anns` calls would exhaust the `group` iterator, thus even if the condition was true, the following code would find no skeletons.
cvat-ai · Oct 14, 2022 · 83c71d2 · 83c71d2
1 parent 0b49dbe
commit 83c71d2
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -35,7 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (<https://github.com/openvinotoolkit/datumaro/pull/712>)
 - Skeleton annotation type
   (<https://github.com/cvat-ai/datumaro/pull/6>,
-  <https://github.com/cvat-ai/datumaro/pull/14>)
+  <https://github.com/cvat-ai/datumaro/pull/14>,
+  <https://github.com/cvat-ai/datumaro/pull/15>)
 - Storing labels with the same name but with a different parent
   (<https://github.com/cvat-ai/datumaro/pull/8>)
 

diff --git a/datumaro/plugins/coco_format/converter.py b/datumaro/plugins/coco_format/converter.py
@@ -397,6 +397,7 @@ def find_solitary_points(cls, annotations):
         solitary_points = []
 
         for g_id, group in groupby(annotations, lambda a: a.group):
+            group = list(group)  # we'll need to iterate over these multiple times
             if not g_id or g_id and not cls.find_instance_anns(group):
                 group = [a for a in group if a.type == AnnotationType.skeleton]
                 solitary_points.extend(group)

diff --git a/tests/test_coco_format.py b/tests/test_coco_format.py
@@ -1717,13 +1717,17 @@ def test_can_save_and_load_keypoints(self):
                         # Full instance annotations: bbox + keypoints
                         Skeleton([Points([1, 2]), Points([3, 4]), Points([2, 3])], group=2, id=2),
                         Bbox(1, 2, 2, 2, group=2, id=2),
-                        # Solitary keypoints
+                        # Solitary keypoints without group
                         Skeleton([Points([1, 2]), Points([0, 2]), Points([4, 1])], label=5, id=3),
+                        # Solitary keypoints with group
+                        Skeleton(
+                            [Points([1, 2]), Points([2, 1]), Points([3, 4])], label=4, group=3, id=4
+                        ),
                         # Some other solitary annotations (bug #1387)
-                        Polygon([0, 0, 4, 0, 4, 4], label=3, id=4),
+                        Polygon([0, 0, 4, 0, 4, 4], label=3, id=5),
                         # Solitary keypoints with no label
                         Skeleton(
-                            [Points([0, 0], [0]), Points([1, 2], [1]), Points([3, 4], [2])], id=5
+                            [Points([0, 0], [0]), Points([1, 2], [1]), Points([3, 4], [2])], id=6
                         ),
                     ],
                 ),
@@ -1772,13 +1776,21 @@ def test_can_save_and_load_keypoints(self):
                             attributes={"is_crowd": False},
                         ),
                         Bbox(0, 1, 4, 1, label=5, group=3, id=3, attributes={"is_crowd": False}),
+                        Skeleton(
+                            [Points([1, 2]), Points([2, 1]), Points([3, 4])],
+                            label=4,
+                            group=4,
+                            id=4,
+                            attributes={"is_crowd": False},
+                        ),
+                        Bbox(1, 1, 2, 3, label=4, group=4, id=4, attributes={"is_crowd": False}),
                         Skeleton(
                             [Points([0, 0], [0]), Points([1, 2], [1]), Points([3, 4], [2])],
-                            group=5,
-                            id=5,
+                            group=6,
+                            id=6,
                             attributes={"is_crowd": False},
                         ),
-                        Bbox(1, 2, 2, 2, group=5, id=5, attributes={"is_crowd": False}),
+                        Bbox(1, 2, 2, 2, group=6, id=6, attributes={"is_crowd": False}),
                     ],
                     attributes={"id": 1},
                 ),