diff --git a/datumaro/datumaro/components/converters/ms_coco.py b/datumaro/datumaro/components/converters/ms_coco.py index 80905771a6e..ed004baa3d6 100644 --- 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 index 9e5701bdc84..5abc9d09c85 100644 --- 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 index 292c2e1c969..9bea8a1f2bc 100644 --- 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)