diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index a01a68b3..08da424b 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -58,34 +58,29 @@ def _polygons_to_masks( def coco_annotations_to_detections( image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool ) -> Detections: - detection = Detections.empty() - class_ids = [] - xyxy = [] - polygons = [] - - for image_annotation in image_annotations: - bbox = image_annotation["bbox"] - xyxy.append(bbox) - class_ids.append(image_annotation["category_id"]) - if with_masks: - _polygons = image_annotation["segmentation"] - _polygons = np.asarray(_polygons, dtype=np.int32) - _polygons = np.reshape(_polygons, (-1, 2)) - polygons.append(_polygons) + if not image_annotations: + return Detections.empty() + class_ids = [ + image_annotation["category_id"] for image_annotation in image_annotations + ] + xyxy = [image_annotation["bbox"] for image_annotation in image_annotations] xyxy = np.asarray(xyxy) - if xyxy.shape[0] > 0: - xyxy[:, 2] += xyxy[:, 0] - xyxy[:, 3] += xyxy[:, 1] - class_ids = np.asarray(class_ids, dtype=int) + xyxy[:, 2:4] += xyxy[:, 0:2] - if with_masks: - mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh) - detection = Detections(class_id=class_ids, xyxy=xyxy, mask=mask) - else: - detection = Detections(xyxy=xyxy, class_id=class_ids) + if with_masks: + polygons = [ + np.reshape( + np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2) + ) + for image_annotation in image_annotations + ] + mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh) + return Detections( + class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask + ) - return detection + return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int)) def detections_to_coco_annotations( @@ -97,7 +92,8 @@ def detections_to_coco_annotations( approximation_percentage: float = 0.75, ) -> Tuple[List[Dict], int]: coco_annotations = [] - for xyxy, mask, confidence, class_id, tracker_id in detections: + for xyxy, mask, _, class_id, _ in detections: + box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1] polygon = [] if mask is not None: polygon = list( @@ -108,16 +104,15 @@ def detections_to_coco_annotations( approximation_percentage=approximation_percentage, )[0].flatten() ) - - coco_annotation = {} - coco_annotation["id"] = annotation_id - coco_annotation["image_id"] = image_id - coco_annotation["category_id"] = int(class_id) - box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1] - coco_annotation["bbox"] = [xyxy[0], xyxy[1], box_width, box_height] - coco_annotation["area"] = box_width * box_height - coco_annotation["segmentation"] = polygon - coco_annotation["iscrowd"] = 0 + coco_annotation = { + "id": annotation_id, + "image_id": image_id, + "category_id": int(class_id), + "bbox": [xyxy[0], xyxy[1], box_width, box_height], + "area": box_width * box_height, + "segmentation": polygon, + "iscrowd": 0, + } coco_annotations.append(coco_annotation) annotation_id += 1 return coco_annotations, annotation_id @@ -128,17 +123,6 @@ def load_coco_annotations( annotations_path: str, force_masks: bool = False, ) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: - """ - Loads COCO annotations and returns class names, images, and their corresponding detections. - - Args: - images_directory_path (str): The path to the directory containing the images. - annotations_path (str): The path to the coco json annotation file. - force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. - - 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. - """ coco_data = read_json_file(file_path=annotations_path) classes = coco_categories_to_classes(coco_categories=coco_data["categories"]) coco_images = coco_data["images"] diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py index 0a75d2aa..3b8d8856 100644 --- a/test/dataset/formats/test_coco.py +++ b/test/dataset/formats/test_coco.py @@ -3,8 +3,11 @@ from typing import List, Tuple 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 + group_coco_annotations_by_image_id, coco_annotations_to_detections + +import numpy as np def generate_cock_coco_annotation( @@ -223,3 +226,65 @@ def test_group_coco_annotations_by_image_id( with exception: result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations) assert result == expected_result + + +@pytest.mark.parametrize( + "image_annotations, resolution_wh, with_masks, expected_result, exception", + [ + ( + [], + (1000, 1000), + False, + Detections.empty(), + DoesNotRaise() + ), # empty image annotations + ( + [ + generate_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100) + ], + (1000, 1000), + False, + Detections( + xyxy=np.array([ + [ 0, 0, 100, 100] + ], dtype=np.float32), + class_id=np.array([ + 0 + ], dtype=int) + ), + DoesNotRaise() + ), # single image annotations + ( + [ + generate_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100), + generate_cock_coco_annotation(category_id=0, bbox=(100, 100, 100, 100), area=100 * 100), + ], + (1000, 1000), + False, + Detections( + xyxy=np.array([ + [ 0, 0, 100, 100], + [ 100, 100, 200, 200] + ], dtype=np.float32), + class_id=np.array([ + 0, 0 + ], dtype=int) + ), + DoesNotRaise() + ), # two image annotations + ] +) +def test_coco_annotations_to_detections( + image_annotations: List[dict], + resolution_wh: Tuple[int, int], + with_masks: bool, + expected_result: Detections, + exception: Exception +) -> None: + with exception: + result = coco_annotations_to_detections( + image_annotations=image_annotations, + resolution_wh=resolution_wh, + with_masks=with_masks + ) + assert result == expected_result