From b25bf6c2163de1198022f41c8ec133a8dca95de3 Mon Sep 17 00:00:00 2001 From: fcakyon <34196005+fcakyon@users.noreply.github.com> Date: Sat, 1 May 2021 14:11:16 +0300 Subject: [PATCH] fix coco subsampling and category updating, increase test coverage (#64) * fix get_subsampled_coco * fix update_categories * update tests for new image_dir property of Coco --- sahi/utils/coco.py | 19 +++++++++++++------ tests/test_cocoutils.py | 29 ++++++++++++++++++++++++++--- 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/sahi/utils/coco.py b/sahi/utils/coco.py index ee05b9eb..69b3df53 100644 --- a/sahi/utils/coco.py +++ b/sahi/utils/coco.py @@ -802,7 +802,7 @@ def add_image(self, image): self.images.append(image) - def update_categories(self, desired_name2id, update_image_filepaths=False): + def update_categories(self, desired_name2id, update_image_filenames=False): """ Rearranges category mapping of given COCO object based on given desired_name2id. Can also be used to filter some of the categories. @@ -810,8 +810,8 @@ def update_categories(self, desired_name2id, update_image_filepaths=False): Args: desired_name2id: dict {"big_vehicle": 1, "car": 2, "human": 3} - update_image_filepaths: bool - If True, updates image file_paths with absolute file paths. + update_image_filenames: bool + If True, updates coco image file_names with absolute file paths. """ # init vars currentid2desiredid_mapping = {} @@ -844,8 +844,11 @@ def update_categories(self, desired_name2id, update_image_filepaths=False): # add updated images & annotations for coco_image in copy.deepcopy(self.images): updated_coco_image = CocoImage.from_coco_image_dict(coco_image.json) - if update_image_filepaths: + # update filename to abspath + file_name_is_abspath = True if os.path.abspath(coco_image.file_name) == coco_image.file_name else False + if update_image_filenames and not file_name_is_abspath: updated_coco_image.file_name = str(Path(os.path.abspath(self.image_dir)) / coco_image.file_name) + # update annotations for coco_annotation in coco_image.annotations: current_category_id = coco_annotation.category_id desired_category_id = currentid2desiredid_mapping[current_category_id] @@ -896,7 +899,7 @@ def merge(self, coco, desired_name2id=None, verbose=1): # update categories and image paths for coco in [coco1, coco2]: - coco.update_categories(desired_name2id=desired_name2id, update_image_filepaths=True) + coco.update_categories(desired_name2id=desired_name2id, update_image_filenames=True) # combine images and categories coco1.images.extend(coco2.images) @@ -1131,7 +1134,11 @@ def get_subsampled_coco(self, subsample_ratio=10): Returns: subsampled_coco: sahi.utils.coco.Coco """ - subsampled_coco = Coco(name=self.name, remapping_dict=self.remapping_dict) + subsampled_coco = Coco( + name=self.name, + image_dir=self.image_dir, + remapping_dict=self.remapping_dict + ) subsampled_coco.add_categories_from_coco_category_list(self.json_categories) for image_ind in tqdm(range(0, len(self.images), subsample_ratio)): subsampled_coco.add_image(self.images[image_ind]) diff --git a/tests/test_cocoutils.py b/tests/test_cocoutils.py index 8e5682fe..0eae6dba 100644 --- a/tests/test_cocoutils.py +++ b/tests/test_cocoutils.py @@ -291,17 +291,20 @@ def test_split_coco_as_train_val(self): from sahi.utils.coco import Coco coco_dict_path = "tests/data/coco_utils/combined_coco.json" - coco = Coco.from_coco_dict_or_path(coco_dict_path) + image_dir = "tests/data/coco_utils/" + coco = Coco.from_coco_dict_or_path(coco_dict_path, image_dir=image_dir) result = coco.split_coco_as_train_val( train_split_rate=0.5, numpy_seed=0 ) self.assertEqual(len(result["train_coco"].json["images"]), 1) self.assertEqual(len(result["train_coco"].json["annotations"]), 5) self.assertEqual(result["train_coco"].json["images"][0]["height"], 682) + self.assertEqual(result["train_coco"].image_dir, image_dir) self.assertEqual(len(result["val_coco"].json["images"]), 1) self.assertEqual(len(result["val_coco"].json["annotations"]), 7) self.assertEqual(result["val_coco"].json["images"][0]["height"], 1365) + self.assertEqual(result["val_coco"].image_dir, image_dir) def test_coco2yolo(self): from sahi.utils.coco import Coco @@ -352,7 +355,8 @@ def test_coco_update_categories(self): from sahi.utils.coco import Coco coco_path = "tests/data/coco_utils/terrain2_coco.json" - coco = Coco.from_coco_dict_or_path(coco_path) + image_dir = "tests/data/coco_utils/" + coco = Coco.from_coco_dict_or_path(coco_path, image_dir=image_dir) self.assertEqual(len(coco.json["annotations"]), 5) self.assertEqual(len(coco.json["images"]), 1) @@ -362,6 +366,7 @@ def test_coco_update_categories(self): [{"id": 1, "name": "car", "supercategory": "car"}], ) self.assertEqual(coco.json["annotations"][1]["category_id"], 1) + self.assertEqual(coco.image_dir, image_dir) # update categories desired_name2id = {"human": 1, "car": 2, "big_vehicle": 3} @@ -379,6 +384,7 @@ def test_coco_update_categories(self): ], ) self.assertEqual(coco.json["annotations"][1]["category_id"], 2) + self.assertEqual(coco.image_dir, image_dir) def test_get_imageid2annotationlist_mapping(self): from sahi.utils.coco import get_imageid2annotationlist_mapping @@ -533,13 +539,22 @@ def test_coco_merge(self): coco1.json["annotations"][9]["image_id"], 2, ) + self.assertEqual( + coco1.image_dir, + image_dir, + ) + self.assertEqual( + coco2.image_dir, + image_dir, + ) def test_get_subsampled_coco(self): from sahi.utils.coco import Coco from sahi.utils.file import load_json coco_path = "tests/data/coco_utils/visdrone2019-det-train-first50image.json" - coco = Coco.from_coco_dict_or_path(coco_path) + image_dir = "tests/data/coco_utils/" + coco = Coco.from_coco_dict_or_path(coco_path, image_dir=image_dir) subsampled_coco = coco.get_subsampled_coco(subsample_ratio=5) self.assertEqual( len(coco.json["images"]), @@ -557,6 +572,14 @@ def test_get_subsampled_coco(self): len(coco.images[5].annotations), len(subsampled_coco.images[1].annotations), ) + self.assertEqual( + coco.image_dir, + image_dir, + ) + self.assertEqual( + subsampled_coco.image_dir, + image_dir, + ) def test_cocovid(self): from sahi.utils.coco import CocoVid