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_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 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_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 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_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)
def test_merge_datasets(): datasets = [ CocoDataset(coco_path=BASE_PATH / dataset_name) for dataset_name in ['dataset1.json', 'dataset2.json'] ] merged_ds = merge_datasets(datasets, BASE_PATH / 'fake_merge.json') assert len(merged_ds.anns) == 20 for k in ['scratch', 'dent']: assert merged_ds.count_images_per_category()[k] == np.sum( [ds_item.count_images_per_category()[k] for ds_item in datasets]) assert merged_ds.count_annotations_per_category()[k] == np.sum([ ds_item.count_annotations_per_category()[k] for ds_item in datasets ]) # check image meta/ anns consistency for ds_item in datasets: for img_idx, img_meta in ds_item.imgs.items(): anns = ds_item.get_annotations(img_idx) merged_img_meta = get_img_meta_by_name(merged_ds, img_meta['file_name']) merged_img_idx = merged_img_meta['id'] merged_anns = merged_ds.get_annotations(merged_img_idx) # number of ann per image assert len(merged_anns) == len(anns) # number of cats per image merged_cats = {a['category_id'] for a in merged_anns} cats = {a['category_id'] for a in anns} assert len(merged_cats) == len(cats) # number of annotated points per image merged_segm_points = {len(a['segmentation']) for a in merged_anns} segm_points = {len(a['segmentation']) for a in anns} assert len(merged_segm_points) == len(segm_points) # annotation areas per image merged_areas = {a['area'] for a in merged_anns} areas = {a['area'] for a in anns} assert merged_areas == areas
def test_load_coco(dataset_file): coco = CocoDataset(dataset_file) assert len(coco.cats) == 3
def coco_test(dataset_file): return CocoDataset(dataset_file)
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'
def test_copy(coco_test: CocoDataset): copy_coco = coco_test.copy() assert id(coco_test.imgs) != id(copy_coco.imgs) assert id(coco_test.cats) != id(copy_coco.cats) assert id(coco_test.anns) != id(copy_coco.anns)
def test_create_split_notest(coco_test: CocoDataset): train, val, test = coco_test.split(.8, .3) assert len(train) > 0 assert len(val) > 0 assert len(test) == 0