Skip to content

Commit

Permalink
[Datumaro] Fix coco images export (#875)
Browse files Browse the repository at this point in the history
* Update test
* Fix export
* Support several image paths in coco extractor
  • Loading branch information
zhiltsov-max authored and nmanovic committed Nov 28, 2019
1 parent 50c40ba commit c3d464c
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 10 deletions.
6 changes: 3 additions & 3 deletions datumaro/datumaro/components/converters/ms_coco.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down
13 changes: 9 additions & 4 deletions datumaro/datumaro/components/extractors/ms_coco.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 7 additions & 3 deletions datumaro/tests/test_coco_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)
Expand Down

0 comments on commit c3d464c

Please sign in to comment.