[Datumaro] Fix coco images export (cvat-ai#875)

* Update test * Fix export * Support several image paths in coco extractor
cleemesser · Mar 5, 2020 · 058745c · 058745c
1 parent 625f20c
commit 058745c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/datumaro/datumaro/components/converters/ms_coco.py b/datumaro/datumaro/components/converters/ms_coco.py
@@ -372,8 +372,8 @@ def make_task_converters(self):
             task: self.make_task_converter(task) for task in self._task
         }
 
-    def save_image(self, item, subset_name, filename):
-        path = osp.join(self._images_dir, subset_name, filename)
+    def save_image(self, item, filename):
+        path = osp.join(self._images_dir, filename)
         cv2.imwrite(path, item.image)
 
         return path
@@ -400,7 +400,7 @@ def convert(self):
                 if item.has_image:
                     filename = str(item.id) + CocoPath.IMAGE_EXT
                     if self._save_images:
-                        self.save_image(item, subset_name, filename)
+                        self.save_image(item, filename)
                 for task_conv in task_converters.values():
                     task_conv.save_image_info(item, filename)
                     task_conv.save_annotations(item)

diff --git a/datumaro/datumaro/components/extractors/ms_coco.py b/datumaro/datumaro/components/extractors/ms_coco.py
@@ -177,10 +177,15 @@ def _get(self, img_id, subset):
                 image_info = loader.loadImgs(img_id)[0]
                 file_name = image_info['file_name']
                 if file_name != '':
-                    file_path = osp.join(
-                        self._path, CocoPath.IMAGES_DIR, subset, file_name)
-                    if osp.exists(file_path):
-                        image = lazy_image(file_path)
+                    image_dir = osp.join(self._path, CocoPath.IMAGES_DIR)
+                    search_paths = [
+                        osp.join(image_dir, file_name),
+                        osp.join(image_dir, subset, file_name),
+                    ]
+                    for image_path in search_paths:
+                        if osp.exists(image_path):
+                            image = lazy_image(image_path)
+                            break
 
             annIds = loader.getAnnIds(imgIds=img_id)
             anns = loader.loadAnns(annIds)

diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py
@@ -101,10 +101,13 @@ def COCO_dataset_generate(self, path):
         ann_dir = osp.join(path, 'annotations')
         os.makedirs(img_dir)
         os.makedirs(ann_dir)
-        a = np.random.rand(100, 100, 3) * 255
-        im_out = Image.fromarray(a.astype('uint8')).convert('RGB')
-        im_out.save(osp.join(img_dir, '000000000001.jpg'))
+
+        image = np.ones((10, 5, 3), dtype=np.uint8)
+        image = Image.fromarray(image).convert('RGB')
+        image.save(osp.join(img_dir, '000000000001.jpg'))
+
         annotation = self.generate_annotation()
+
         with open(osp.join(ann_dir, 'instances_val.json'), 'w') as outfile:
             json.dump(annotation, outfile)
 
@@ -119,6 +122,7 @@ def test_can_import(self):
 
             item = next(iter(dataset))
             self.assertTrue(item.has_image)
+            self.assertEqual(np.sum(item.image), np.prod(item.image.shape))
             self.assertEqual(4, len(item.annotations))
 
             ann_1 = find(item.annotations, lambda x: x.id == 1)