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

* Update test
* Fix export
* Support several image paths in coco extractor

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)

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)

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)