From 6ca9555efd608082044b6e7dec65dbce4bd92e16 Mon Sep 17 00:00:00 2001 From: Youho99 <44434482+Youho99@users.noreply.github.com> Date: Fri, 20 Feb 2026 12:08:26 +0100 Subject: [PATCH] Fix `detections_to_coco_annotations` function for empty polygons. (#1086) * Fix `detections_to_coco_annotations` function for empty polygons. * Add `segmentation` empty field for bboxes coco format * update coco.py Always include box, area, and segmentation in the result COCO JSON. * `as_coco()` : Add COCO format disjoint masks support * Fix `force_masks` parameter on `from_coco()` function Allows reconstruction of disjointed masks * Ensures the mask is binary * Refactor `coco_annotations_to_masks` to handle disjoint polygon segmentation and missing segmentation gracefully. Add corresponding unit tests. * Refactor `coco_annotations_to_masks` for improved type annotations and cleaner formatting * Add warnings for handling empty polygons during COCO segmentation and related unit tests --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: jirka <6035284+Borda@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/supervision/dataset/formats/coco.py | 75 +++++++++++++----- tests/dataset/formats/test_coco.py | 101 ++++++++++++++++++++++++ 2 files changed, 158 insertions(+), 18 deletions(-) diff --git a/src/supervision/dataset/formats/coco.py b/src/supervision/dataset/formats/coco.py index eb2fca22..001f8145 100644 --- a/src/supervision/dataset/formats/coco.py +++ b/src/supervision/dataset/formats/coco.py @@ -1,4 +1,5 @@ import os +import warnings from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any, Union @@ -68,24 +69,52 @@ def group_coco_annotations_by_image_id( def coco_annotations_to_masks( image_annotations: list[CocoDict], resolution_wh: tuple[int, int] ) -> npt.NDArray[np.bool_]: - return np.array( - [ - rle_to_mask( - rle=np.array(image_annotation["segmentation"]["counts"]), - resolution_wh=resolution_wh, + height, width = resolution_wh[1], resolution_wh[0] + empty_mask = np.zeros((height, width), dtype=bool) + masks = [] + + for image_annotation in image_annotations: + segmentation = image_annotation.get("segmentation") + if not segmentation: + # `force_masks=True` may request masks even for bbox-only annotations. + # Keep detection count aligned by emitting an empty mask for that object. + masks.append(empty_mask.copy()) + continue + + if image_annotation.get("iscrowd", 0): + masks.append( + rle_to_mask( + rle=np.array(segmentation["counts"]), resolution_wh=resolution_wh + ).astype(bool) ) - if image_annotation["iscrowd"] - else polygon_to_mask( - polygon=np.reshape( - np.asarray(image_annotation["segmentation"], dtype=np.int32), - (-1, 2), - ), - resolution_wh=resolution_wh, + continue + + if not isinstance(segmentation, list): + masks.append(empty_mask.copy()) + continue + polygons = segmentation if isinstance(segmentation[0], list) else [segmentation] + + object_mask = empty_mask.copy() + for polygon in polygons: + polygon_array: npt.NDArray[np.int32] = np.reshape( + np.asarray(polygon, dtype=np.int32), (-1, 2) ) - for image_annotation in image_annotations - ], - dtype=bool, - ) + if polygon_array.size == 0: + warnings.warn( + "Skipping empty polygon while loading COCO segmentation for " + f"annotation id={image_annotation.get('id')}.", + stacklevel=2, + ) + continue + # COCO polygon segmentation can contain multiple disjoint parts. + # Merge all parts into a single per-object mask. + object_mask |= polygon_to_mask( + polygon=polygon_array, resolution_wh=resolution_wh + ).astype(bool) + + masks.append(object_mask) + + return np.asarray(masks, dtype=bool) def coco_annotations_to_detections( @@ -156,7 +185,17 @@ def detections_to_coco_annotations( max_image_area_percentage=max_image_area_percentage, approximation_percentage=approximation_percentage, ) - segmentation = [list(polygons[0].flatten())] + # Small/noisy masks can be filtered out by approximation settings. + # Guard against empty output and keep a valid COCO annotation record. + if polygons: + segmentation = [list(polygons[0].flatten())] + else: + warnings.warn( + "Skipping COCO polygon segmentation for annotation " + f"id={annotation_id} because mask approximation " + "returned no polygons.", + stacklevel=2, + ) coco_annotation = { "id": annotation_id, "image_id": image_id, @@ -278,7 +317,7 @@ def load_coco_annotations( def _with_seg_mask(annotation: dict[str, Any]) -> bool: - return "segmentation" in annotation + return bool(annotation.get("segmentation")) def save_coco_annotations( diff --git a/tests/dataset/formats/test_coco.py b/tests/dataset/formats/test_coco.py index caabdaab..586b3e48 100644 --- a/tests/dataset/formats/test_coco.py +++ b/tests/dataset/formats/test_coco.py @@ -375,6 +375,39 @@ def test_group_coco_annotations_by_image_id( ), DoesNotRaise(), ), # single image annotations with mask as polygon + ( + [ + mock_coco_annotation( + category_id=0, + bbox=(0, 0, 5, 5), + area=5 * 5, + segmentation=[ + [0, 0, 1, 0, 1, 1, 0, 1], + [3, 3, 4, 3, 4, 4, 3, 4], + ], + ) + ], + (5, 5), + True, + False, + Detections( + xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32), + class_id=np.array([0], dtype=int), + mask=np.array( + [ + [ + [1, 1, 0, 0, 0], + [1, 1, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 1, 1], + [0, 0, 0, 1, 1], + ] + ], + dtype=bool, + ), + ), + DoesNotRaise(), + ), # single image annotation with disjoint polygon segments ( [ mock_coco_annotation( @@ -883,6 +916,37 @@ def test_detections_to_coco_annotations( assert result == expected_result +def test_detections_to_coco_annotations_handles_empty_approximated_polygons() -> None: + detections = Detections( + xyxy=np.array([[0, 0, 4, 4]], dtype=np.float32), + class_id=np.array([0], dtype=int), + mask=np.array( + [ + [ + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 0], + ] + ], + dtype=bool, + ), + ) + + with pytest.warns(Warning, match="mask approximation returned no polygons"): + annotations, _ = detections_to_coco_annotations( + detections=detections, + image_id=0, + annotation_id=0, + max_image_area_percentage=0.01, + ) + + assert len(annotations) == 1 + assert annotations[0]["segmentation"] == [] + assert annotations[0]["iscrowd"] == 0 + + def test_load_coco_annotations_infers_masks_from_segmentation_field( tmp_path, coco_data_with_and_without_segmentation: dict[str, object] ) -> None: @@ -954,3 +1018,40 @@ def test_load_coco_annotations_force_masks_with_no_annotations( no_annotations_path = str(images_directory / "no_annotations.jpg") assert annotations[no_annotations_path] == Detections.empty() + + +def test_load_coco_annotations_force_masks_handles_missing_segmentation( + tmp_path, +) -> None: + images_directory = tmp_path / "images" + images_directory.mkdir() + annotations_path = tmp_path / "annotations.json" + + coco_data = { + "categories": [{"id": 1, "name": "object", "supercategory": "none"}], + "images": [{"id": 1, "file_name": "image.jpg", "width": 5, "height": 5}], + "annotations": [ + { + "id": 1, + "image_id": 1, + "category_id": 1, + "bbox": [0, 0, 2, 2], + "area": 4, + "iscrowd": 0, + } + ], + } + annotations_path.write_text(json.dumps(coco_data), encoding="utf-8") + + _, _, annotations = load_coco_annotations( + images_directory_path=str(images_directory), + annotations_path=str(annotations_path), + force_masks=True, + ) + + image_path = str(images_directory / "image.jpg") + image_annotations = annotations[image_path] + assert image_annotations.mask is not None + assert image_annotations.mask.shape == (1, 5, 5) + assert not image_annotations.mask.any() + assert np.array_equal(image_annotations.xyxy, np.array([[0, 0, 2, 2]], dtype=float))