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'