okooooooo/supervision
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix detections_to_coco_annotations function for empty polygons. (#1086)

Browse Source 
* 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>
This commit is contained in:
Youho99pre-commit-ci[bot]jirkaCopilot
 authored and GitHub committed 2026-02-20 12:08:26 +01:00
1 parent ae2c356321
commit 6ca9555efd
2 files changed
+158 -18
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/supervision/dataset/formats/coco.py
+57 -18
View File
@@ -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(
tests/dataset/formats/test_coco.py
+101
View File
@@ -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))