def test_get_annotations_for_image(coco_test: CocoDataset):
coco_test.reindex()
img_idx = 1
ann_idxs = coco_test.index.imgidx_to_annidxs[img_idx]
coco_test.remove_annotations(ann_idxs)
anns = coco_test.get_annotations(img_idx)
assert isinstance(anns, list)
assert len(anns) == 0
def merge_datasets(
datasets: List[Union[SemanticCocoDataset, CocoDataset, InstanceCoco]],
dest_path: Union[str, Path] = None,
) -> CocoDataset:
# compute category to cat_idx map for each dataset
cat_to_idx = [{v["name"]: k
for k, v in ds_item.cats.items()} for ds_item in datasets]
category_names = [list(ctx.keys()) for ctx in cat_to_idx]
flat_category_name = sorted(
[item for sublist in category_names for item in sublist])
# collect all category names
merged_cat_names = set(flat_category_name)
merged_idx_to_cat = dict(enumerate(sorted(merged_cat_names), start=1))
merged_cat_to_idx = {v: k for k, v in merged_idx_to_cat.items()}
# map categories for each dataset to merged one
cat_map = [{v: merged_cat_to_idx[k]
for k, v in ctx.items()} for ctx in cat_to_idx]
# initializize merged dataset
out_ds = CocoDataset(dest_path)
for _, cat_name in sorted(merged_idx_to_cat.items()):
out_ds.add_category(cat_name, "thing")
for ds_item, cat_map_item in zip(datasets, cat_map):
img_names = {a["file_name"]: a["id"] for a in out_ds.imgs.values()}
# add images
for img_meta in ds_item.imgs.values():
# remark: if datasets are not disjoint the might bo overlap here to take care of, i.e. possible multiple definitions of images
if img_meta["file_name"] not in img_names:
img_idx = out_ds.add_image(**img_meta)
img_names[img_meta["file_name"]] = img_idx
else:
img_idx = img_names[img_meta["file_name"]]
orig_img_idx = img_meta["id"]
anns = ds_item.get_annotations(img_idx=orig_img_idx)
# add anns
for ann in anns:
mapped_ann = rename_annotation_keys(ann)
mapped_ann["img_id"] = img_idx # new image idx
mapped_ann["cat_id"] = cat_map_item[
mapped_ann["cat_id"]] # new category index
out_ds.add_annotation(**mapped_ann)
out_ds.reindex()
return out_ds
def test_get_cvat_categories(coco_test: CocoDataset):
coco_test.reindex()
expected = [{
'name': 'bear',
'color': '',
'attributes': []
}, {
'name': 'toaster',
'color': '',
'attributes': []
}, {
'name': 'hair drier',
'color': '',
'attributes': []
}]
cvat_categories = coco_test.get_cvat_categories()
assert cvat_categories == expected
def test_update_images_path(coco_test: CocoDataset):
coco_test.update_images_path(lambda x: Path(x).name)
coco_test.reindex()
assert coco_test.imgs[1]['file_name'] == '000000001442.jpg'
