diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 7d344de5..091bb19e 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -12,13 +12,23 @@ from supervision.detection.utils import polygon_to_mask from supervision.utils.file import read_json_file, save_json_file -def _extract_class_names(annotation_data: dict) -> List[str]: - names = [] - categories = annotation_data.get("categories", None) - if categories: - for cat in categories: - names.append(cat["name"]) - return names +def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]: + return [ + category["name"] + for category in sorted(coco_categories, key=lambda category: category["id"]) + if category["supercategory"] != "none" + ] + + +def classes_to_coco_categories(classes: List[str]) -> List[dict]: + return [ + { + "id": class_id, + "name": class_name, + "supercategory": "common-objects", + } + for class_id, class_name in enumerate(classes) + ] def _extract_image_info(annotation_data: dict) -> List[dict]: @@ -26,13 +36,6 @@ def _extract_image_info(annotation_data: dict) -> List[dict]: return image_infos -def _extract_image_names(image_infos: dict) -> List[dict]: - image_names = [] - for image_info in image_infos: - image_names.append(image_info["file_name"]) - return image_names - - def _annotations_dict(annotation_data: dict) -> Tuple[np.ndarray, dict]: annotations_infos = annotation_data.get("annotations", None) image_id_label_id_pair = np.zeros((0, 2)) @@ -128,18 +131,6 @@ def detections_to_coco_annotations( return annotations_data, label_id -def classes_to_coco_category_map(classes: List[str]) -> List[Dict]: - categories = [] - for class_id, class_name in enumerate(classes): - cate_dict = { - "id": class_id, - "name": class_name, - "supercategory": "common-objects", - } - categories.append(cate_dict) - return categories - - def load_coco_annotations( images_directory_path: str, annotations_path: str, @@ -158,7 +149,7 @@ def load_coco_annotations( """ annotation_data = read_json_file(file_path=annotations_path) - classes = _extract_class_names(annotation_data=annotation_data) + classes = coco_categories_to_classes(coco_categories=annotation_data["categories"]) images_infos = _extract_image_info(annotation_data=annotation_data) image_id_label_id_pair, annotation_dict = _annotations_dict( annotation_data=annotation_data @@ -168,7 +159,8 @@ def load_coco_annotations( annotations = {} for images_info in images_infos: - image_path = os.path.join(images_directory_path, images_info["file_name"]) + image_name = images_info["file_name"] + image_path = os.path.join(images_directory_path, image_name) image = cv2.imread(str(image_path)) # Filter annotations based on image id @@ -188,8 +180,8 @@ def load_coco_annotations( with_masks=force_masks, ) - images[images_info["file_name"]] = image - annotations[images_info["file_name"]] = annotation + images[image_name] = image + annotations[image_name] = annotation return classes, images, annotations @@ -219,7 +211,7 @@ def save_coco_annotations( annotations_data = [] image_infos = [] - categories = classes_to_coco_category_map(classes=classes) + categories = classes_to_coco_categories(classes=classes) image_id = 0 label_id = 0 diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py new file mode 100644 index 00000000..f0ed80b5 --- /dev/null +++ b/test/dataset/formats/test_coco.py @@ -0,0 +1,132 @@ +from contextlib import ExitStack as DoesNotRaise +from typing import List + +import pytest + +from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes + + +@pytest.mark.parametrize( + "coco_categories, expected_result, exception", + [ + ( + [], + [], + DoesNotRaise() + ), # empty coco categories + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + } + ], + [], + DoesNotRaise() + ), # single coco category with supercategory == "none" + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap" + ], + DoesNotRaise() + ), # two coco categories; one with supercategory == "none" and one with supercategory != "none" + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + }, + { + "id": 2, + "name": "hoodie", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap", + "hoodie" + ], + DoesNotRaise() + ), # three coco categories; one with supercategory == "none" and two with supercategory != "none" + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + }, + { + "id": 2, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap", + "hoodie" + ], + DoesNotRaise() + ), # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order) + ] +) +def test_coco_categories_to_classes( + coco_categories: List[dict], + expected_result: List[str], + exception: Exception +) -> None: + with exception: + result = coco_categories_to_classes(coco_categories=coco_categories) + assert result == expected_result + + +@pytest.mark.parametrize( + "classes, exception", + [ + ( + [], + DoesNotRaise() + ), # empty classes + ( + [ + "baseball cap" + ], + DoesNotRaise() + ), # single class + ( + [ + "baseball cap", + "hoodie" + ], + DoesNotRaise() + ), # two classes + ] +) +def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exception: Exception) -> None: + with exception: + coco_categories = classes_to_coco_categories(classes=classes) + result = coco_categories_to_classes(coco_categories=coco_categories) + assert result == classes