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_remove_categories():
ds = CocoDataset(BASE_PATH / 'new_coco.json')
cat_id = ds.add_category('dog', 'animal')
assert len(ds.cats) == 1
ds.remove_categories([cat_id])
assert len(ds.cats) == 0
def test_remove_categories():
ds = CocoDataset(BASE_PATH / "new_coco.json")
cat_id = ds.add_category("dog", "animal")
assert len(ds.cats) == 1
ds.remove_categories([cat_id])
assert len(ds.cats) == 0
def test_remove_categories_and_annotations():
ds = CocoDataset(BASE_PATH / 'new_coco.json')
cat_id = ds.add_category('dog', 'animal')
img_id = ds.add_image((BASE_PATH / 'test_nodamage.jpg').as_posix())
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
assert len(ds.cats) == 1
ds.remove_categories([cat_id])
assert len(ds.cats) == 0
def test_remove_categories_and_annotations():
ds = CocoDataset(BASE_PATH / "new_coco.json")
cat_id = ds.add_category("dog", "animal")
img_id = ds.add_image((BASE_PATH / "test_nodamage.jpg").as_posix(), 100,
100)
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
assert len(ds.cats) == 1
ds.remove_categories([cat_id])
assert len(ds.cats) == 0
def test_create_dataset_existing():
ds = CocoDataset(BASE_PATH / 'new_coco.json')
cat_id = ds.add_category('dog', 'animal')
img_id = ds.add_image((BASE_PATH / 'test_nodamage.jpg').as_posix())
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
img_id = ds.add_image((BASE_PATH / 'test_nodamage.jpg').as_posix())
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
assert len(ds.imgs) == 2
assert len(ds.anns) == 2
assert ds.cat_id == 2
assert ds.img_id == 2
assert ds.ann_id == 2
def test_create_dataset_existing():
ds = CocoDataset(BASE_PATH / "new_coco.json")
cat_id = ds.add_category("dog", "animal")
img_id = ds.add_image((BASE_PATH / "test_nodamage.jpg").as_posix(), 100,
100)
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
img_id = ds.add_image((BASE_PATH / "test_nodamage.jpg").as_posix(), 100,
100)
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
assert len(ds.imgs) == 2
assert len(ds.anns) == 2
assert ds.cat_id == 2
assert ds.img_id == 3
assert ds.ann_id == 3
def test_create_dataset():
annotations_path = BASE_PATH / 'new_coco.json'
ds = CocoDataset(annotations_path)
cat_id = ds.add_category('dog', 'animal')
img_id = ds.add_image((BASE_PATH / 'test_nodamage.jpg').as_posix())
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
assert len(ds) == 1
assert len(ds.anns) == 1
ds.dump()
assert annotations_path.exists()
with open(annotations_path, 'r') as f:
assert len(f.readline()) > 0
os.remove(annotations_path)
def test_create_dataset():
annotations_path = BASE_PATH / "new_coco.json"
ds = CocoDataset(annotations_path)
cat_id = ds.add_category("dog", "animal")
img_id = ds.add_image((BASE_PATH / "test_nodamage.jpg").as_posix(), 100,
100)
ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
assert len(ds) == 1
assert len(ds.anns) == 1
ds.dump()
assert annotations_path.exists()
with open(annotations_path, "r") as f:
assert len(f.readline()) > 0
os.remove(annotations_path)