Fix COCO Panoptic (cvat-ai#319)

* add test * Fix integer overflow in bgr2index
TOsmanov · Jul 5, 2021 · f955177 · f955177
1 parent 37b8cb0
commit f955177
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 7 deletions.
diff --git a/datumaro/util/mask_tools.py b/datumaro/util/mask_tools.py
@@ -114,8 +114,8 @@ def make_binary_mask(mask):
     return mask.astype(bool)
 
 def bgr2index(img):
-    if img.dtype.kind not in {'b', 'i', 'u'}:
-        img = img.astype(np.uint8)
+    if img.dtype.kind not in {'b', 'i', 'u'} or img.dtype.itemsize < 4:
+        img = img.astype(np.uint32)
     return (img[..., 0] << 16) + (img[..., 1] << 8) + img[..., 2]
 
 def index2bgr(id_map):

diff --git a/tests/assets/coco_dataset/coco_panoptic/annotations/panoptic_train.json b/tests/assets/coco_dataset/coco_panoptic/annotations/panoptic_train.json
@@ -46,7 +46,7 @@
         "file_name":"a.png",
         "segments_info":[
           {
-            "id":7,
+            "id":460551,
             "category_id":1,
             "area":20.0,
             "bbox":[

diff --git a/tests/assets/coco_dataset/coco_panoptic/annotations/panoptic_train/a.png b/tests/assets/coco_dataset/coco_panoptic/annotations/panoptic_train/a.png
diff --git a/tests/test_coco_format.py b/tests/test_coco_format.py
@@ -230,10 +230,8 @@ def test_can_import_panoptic(self):
             DatasetItem(id='a', subset='train', image=np.ones((5, 10, 3)),
                 attributes={'id': 5},
                 annotations=[
-                    Mask(np.array(
-                        [[0, 0, 1, 1, 0, 1, 1, 0, 0, 0]] * 5
-                        ), label=0,
-                        id=7, group=7, attributes={'is_crowd': False}),
+                    Mask(np.ones((5, 5)), label=0, id=460551,
+                        group=460551, attributes={'is_crowd': False}),
                 ]),
 
             DatasetItem(id='b', subset='val', image=np.ones((10, 5, 3)),