diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 1c0e6b55..b9ea4fdb 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -31,24 +31,16 @@ def classes_to_coco_categories(classes: List[str]) -> List[dict]: ] -def _extract_image_info(annotation_data: dict) -> List[dict]: - image_infos = annotation_data.get("images", None) - return image_infos - - -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)) +def group_coco_annotations_by_image_id( + coco_annotations: List[dict], +) -> Dict[int, List[dict]]: annotations = {} - for anno in annotations_infos: - _annotations = {} - for key, item in anno.items(): - _annotations[key] = anno[key] - id_pair = np.array([[anno["image_id"], anno["id"]]]) - image_id_label_id_pair = np.append(image_id_label_id_pair, id_pair, axis=0) - annotations[anno["id"]] = _annotations - - return image_id_label_id_pair, annotations + for annotation in coco_annotations: + image_id = annotation["image_id"] + if image_id not in annotations: + annotations[image_id] = [] + annotations[image_id].append(annotation) + return annotations def _polygons_to_masks( @@ -147,36 +139,29 @@ def load_coco_annotations( Returns: Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values. """ - annotation_data = read_json_file(file_path=annotations_path) - - classes = coco_categories_to_classes(coco_categories=annotation_data["categories"]) - coco_images = annotation_data["images"] - image_id_label_id_pair, annotation_dict = _annotations_dict( - annotation_data=annotation_data + coco_data = read_json_file(file_path=annotations_path) + classes = coco_categories_to_classes(coco_categories=coco_data["categories"]) + coco_images = coco_data["images"] + coco_annotations_groups = group_coco_annotations_by_image_id( + coco_annotations=coco_data["annotations"] ) images = {} annotations = {} for coco_image in coco_images: - image_name = coco_image["file_name"] + image_name, image_width, image_height = ( + coco_image["file_name"], + coco_image["width"], + coco_image["height"], + ) + image_annotations = coco_annotations_groups.get(coco_image["id"]) image_path = os.path.join(images_directory_path, image_name) + image = cv2.imread(str(image_path)) - - # Filter annotations based on image id - per_image_label_ids = image_id_label_id_pair[ - image_id_label_id_pair[:, 0] == coco_image["id"] - ][:, 1] - - image_annotations = [] - for per_image_label_id in per_image_label_ids: - image_annotations.append(annotation_dict[int(per_image_label_id)]) - - w, h = coco_image["width"], coco_image["height"] - annotation = coco_annotations_to_detections( image_annotations=image_annotations, - resolution_wh=(w, h), + resolution_wh=(image_width, image_height), with_masks=force_masks, ) diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py index f0ed80b5..0a75d2aa 100644 --- a/test/dataset/formats/test_coco.py +++ b/test/dataset/formats/test_coco.py @@ -1,9 +1,27 @@ from contextlib import ExitStack as DoesNotRaise -from typing import List +from typing import List, Tuple import pytest -from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes +from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \ + group_coco_annotations_by_image_id + + +def generate_cock_coco_annotation( + annotation_id: int = 0, + image_id: int = 0, + category_id: int = 0, + bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + area: float = 0.0 +) -> dict: + return { + "id": annotation_id, + "image_id": image_id, + "category_id": category_id, + "bbox": list(bbox), + "area": area, + "iscrowd": 0 + } @pytest.mark.parametrize( @@ -130,3 +148,78 @@ def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exce coco_categories = classes_to_coco_categories(classes=classes) result = coco_categories_to_classes(coco_categories=coco_categories) assert result == classes + + +@pytest.mark.parametrize( + "coco_annotations, expected_result, exception", + [ + ( + [], + {}, + DoesNotRaise() + ), # empty coco annotations + ( + [ + generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0) + ], + { + 0: [ + generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0) + ] + }, + DoesNotRaise() + ), # single coco annotation + ( + [ + generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0), + generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0) + ], + { + 0: [ + generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0) + ], + 1: [ + generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0) + ] + }, + DoesNotRaise() + ), # two coco annotations + ( + [ + generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0), + generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1), + generate_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2), + generate_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3), + generate_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1), + generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2), + generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3), + ], + { + 0: [ + generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0), + ], + 1: [ + generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1), + generate_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2), + ], + 2: [ + generate_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3), + ], + 3: [ + generate_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1), + generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2), + generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3), + ] + }, + DoesNotRaise() + ), # two coco annotations + ] +) +def test_group_coco_annotations_by_image_id( + coco_annotations: List[dict], + expected_result: dict, + exception: Exception +) -> None: + with exception: + result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations) + assert result == expected_result