diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 51adff01..c1d0a77a 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -6,7 +6,10 @@ from typing import Dict, List, Tuple import cv2 import numpy as np -from supervision.dataset.utils import approximate_mask_with_polygons +from supervision.dataset.utils import ( + approximate_mask_with_polygons, + map_detections_class_id, +) from supervision.detection.core import Detections from supervision.detection.utils import polygon_to_mask from supervision.utils.file import read_json_file, save_json_file @@ -19,6 +22,18 @@ def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]: ] +def build_coco_class_index_mapping( + coco_categories: List[dict], target_classes: List[str] +) -> Dict[int, int]: + source_class_to_index = { + category["name"]: category["id"] for category in coco_categories + } + return { + source_class_to_index[target_class_name]: target_class_index + for target_class_index, target_class_name in enumerate(target_classes) + } + + def classes_to_coco_categories(classes: List[str]) -> List[dict]: return [ { @@ -124,6 +139,9 @@ def load_coco_annotations( ) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: coco_data = read_json_file(file_path=annotations_path) classes = coco_categories_to_classes(coco_categories=coco_data["categories"]) + class_index_mapping = build_coco_class_index_mapping( + coco_categories=coco_data["categories"], target_classes=classes + ) coco_images = coco_data["images"] coco_annotations_groups = group_coco_annotations_by_image_id( coco_annotations=coco_data["annotations"] @@ -147,6 +165,10 @@ def load_coco_annotations( resolution_wh=(image_width, image_height), with_masks=force_masks, ) + annotation = map_detections_class_id( + source_to_target_mapping=class_index_mapping, + detections=annotation, + ) images[image_name] = image annotations[image_name] = annotation diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py index ca1ceb6c..cbbb2e21 100644 --- a/test/dataset/formats/test_coco.py +++ b/test/dataset/formats/test_coco.py @@ -1,11 +1,11 @@ from contextlib import ExitStack as DoesNotRaise -from typing import List, Tuple +from typing import List, Tuple, Dict import pytest from supervision import Detections from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \ - group_coco_annotations_by_image_id, coco_annotations_to_detections + group_coco_annotations_by_image_id, coco_annotations_to_detections, build_coco_class_index_mapping import numpy as np @@ -290,3 +290,135 @@ def test_coco_annotations_to_detections( with_masks=with_masks ) assert result == expected_result + +@pytest.mark.parametrize( + "coco_categories, target_classes, expected_result, exception", + [ + ( + [], + [], + {}, + DoesNotRaise() + ), # empty coco categories + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + } + ], + [ + "fashion-assistant" + ], + { + 0: 0 + }, + DoesNotRaise() + ), # single coco category starting from 0 + ( + [ + { + "id": 1, + "name": "fashion-assistant", + "supercategory": "none" + } + ], + [ + "fashion-assistant" + ], + { + 1: 0 + }, + DoesNotRaise() + ), # single coco category starting from 1 + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + }, + { + "id": 2, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "fashion-assistant", + "baseball cap", + "hoodie" + ], + { + 0: 0, + 1: 1, + 2: 2 + }, + DoesNotRaise() + ), # three coco categories + ( + [ + { + "id": 2, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap", + "hoodie" + ], + { + 2: 1, + 1: 0 + }, + DoesNotRaise() + ), # two coco categories + ( + [ + { + "id": 3, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap", + "hoodie" + ], + { + 3: 1, + 1: 0 + }, + DoesNotRaise() + ), # two coco categories with missing category + ] +) +def test_build_coco_class_index_mapping( + coco_categories: List[dict], + target_classes: List[str], + expected_result: Dict[int, int], + exception: Exception +) -> None: + with exception: + result = build_coco_class_index_mapping( + coco_categories=coco_categories, + target_classes=target_classes + ) + assert result == expected_result