Files
wehub-resource-sync 9194ef5abd
Docs/Test Workflow / Test docs build (push) Failing after 0s
Check links & references / links-check (push) Failing after 1s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.10) (push) Failing after 0s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.11) (push) Failing after 0s
PR Conflict Labeler / main (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.12) (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.13) (push) Failing after 0s
Pytest/Test Workflow / Build this Package (push) Failing after 5s
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / testing-guardian (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:06:10 +08:00

2403 lines
80 KiB
Python

import json
from contextlib import ExitStack as DoesNotRaise
from pathlib import Path
import cv2
import numpy as np
import pytest
from supervision import DetectionDataset, Detections
from supervision.dataset.formats.coco import (
build_coco_class_index_mapping,
classes_to_coco_categories,
coco_annotations_to_detections,
coco_annotations_to_masks,
coco_categories_to_classes,
detections_to_coco_annotations,
group_coco_annotations_by_image_id,
load_coco_annotations,
save_coco_annotations,
)
def mock_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,
segmentation: list[list] | dict | None = None,
iscrowd: bool = False,
) -> dict:
if not segmentation:
segmentation = []
return {
"id": annotation_id,
"image_id": image_id,
"category_id": category_id,
"bbox": list(bbox),
"area": area,
"segmentation": segmentation,
"iscrowd": int(iscrowd),
}
def _empty_raw_segs(n: int) -> np.ndarray:
"""Object-dtype array of n empty lists for coco_raw_segmentation (bbox-only)."""
arr = np.empty(n, dtype=object)
for i in range(n):
arr[i] = []
return arr
@pytest.fixture
def coco_data_with_and_without_segmentation() -> dict[str, object]:
return {
"categories": [{"id": 1, "name": "object", "supercategory": "none"}],
"images": [
{"id": 1, "file_name": "with_segmentation.jpg", "width": 5, "height": 5},
{
"id": 2,
"file_name": "with_polygon_segmentation.jpg",
"width": 5,
"height": 5,
},
{"id": 3, "file_name": "without_segmentation.jpg", "width": 5, "height": 5},
{"id": 4, "file_name": "without_annotations.jpg", "width": 5, "height": 5},
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [0, 0, 5, 5],
"area": 25,
"segmentation": [[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
"iscrowd": 0,
},
{
"id": 2,
"image_id": 1,
"category_id": 1,
"bbox": [3, 0, 2, 2],
"area": 4,
"segmentation": {"size": [5, 5], "counts": [15, 2, 3, 2, 3]},
"iscrowd": 1,
},
{
"id": 3,
"image_id": 2,
"category_id": 1,
"bbox": [0, 0, 2, 2],
"area": 4,
"segmentation": [[0, 0, 1, 0, 1, 1, 0, 1]],
"iscrowd": 0,
},
{
"id": 4,
"image_id": 3,
"category_id": 1,
"bbox": [0, 0, 2, 2],
"area": 4,
"iscrowd": 0,
},
],
}
@pytest.fixture
def coco_data_with_unannotated_image() -> dict[str, object]:
return {
"categories": [{"id": 1, "name": "object", "supercategory": "none"}],
"images": [
{"id": 1, "file_name": "has_segmentation.jpg", "width": 5, "height": 5},
{"id": 2, "file_name": "no_annotations.jpg", "width": 5, "height": 5},
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [0, 0, 2, 2],
"area": 4,
"segmentation": [[0, 0, 1, 0, 1, 1, 0, 1]],
"iscrowd": 0,
}
],
}
@pytest.mark.parametrize(
("coco_categories", "expected_result", "exception"),
[
([], [], DoesNotRaise()), # empty coco categories
(
[{"id": 0, "name": "fashion-assistant", "supercategory": "none"}],
["fashion-assistant"],
DoesNotRaise(),
), # single coco category with supercategory == "none"
(
[
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
],
["fashion-assistant", "baseball cap"],
DoesNotRaise(),
), # two coco categories; one with supercategory == "none" and
# one with supercategory != "none"
(
[
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
],
["fashion-assistant", "baseball cap", "hoodie"],
DoesNotRaise(),
), # three coco categories; one with supercategory == "none" and
# two with supercategory != "none"
(
[
{"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"],
DoesNotRaise(),
), # three coco categories; one with supercategory == "none" and
# two with supercategory != "none" (different order)
],
)
def test_coco_categories_to_classes(
coco_categories: list[dict], expected_result: list[str], exception: Exception
) -> None:
with exception:
result = coco_categories_to_classes(coco_categories=coco_categories)
assert result == expected_result
@pytest.mark.parametrize(
("classes", "exception"),
[
([], DoesNotRaise()), # empty classes
(["baseball cap"], DoesNotRaise()), # single class
(["baseball cap", "hoodie"], DoesNotRaise()), # two classes
],
)
def test_classes_to_coco_categories_and_back_to_classes(
classes: list[str], exception: Exception
) -> None:
with exception:
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
(
[mock_coco_annotation(annotation_id=0, image_id=0, category_id=0)],
{0: [mock_coco_annotation(annotation_id=0, image_id=0, category_id=0)]},
DoesNotRaise(),
), # single coco annotation
(
[
mock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
mock_coco_annotation(annotation_id=1, image_id=1, category_id=0),
],
{
0: [mock_coco_annotation(annotation_id=0, image_id=0, category_id=0)],
1: [mock_coco_annotation(annotation_id=1, image_id=1, category_id=0)],
},
DoesNotRaise(),
), # two coco annotations
(
[
mock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
mock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
mock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
mock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
mock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
mock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
mock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
],
{
0: [
mock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
],
1: [
mock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
mock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
],
2: [
mock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
],
3: [
mock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
mock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
mock_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
def test_coco_annotations_to_detections_defaults_missing_area_and_iscrowd() -> None:
"""COCO annotations missing optional area/iscrowd fields still load."""
annotation = {
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [1, 2, 3, 4],
"segmentation": [],
}
detections = coco_annotations_to_detections(
image_annotations=[annotation],
resolution_wh=(10, 10),
with_masks=False,
use_iscrowd=True,
)
assert detections.data["iscrowd"].tolist() == [0]
assert detections.data["area"].tolist() == [12.0]
def test_load_coco_annotations_can_skip_iscrowd_metadata(tmp_path: Path) -> None:
"""use_iscrowd=False omits COCO metadata from loaded Detections.data."""
images_dir = tmp_path / "images"
images_dir.mkdir()
image_path = images_dir / "image.jpg"
cv2.imwrite(str(image_path), np.zeros((10, 10, 3), dtype=np.uint8))
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(
json.dumps(
{
"categories": [{"id": 1, "name": "object"}],
"images": [
{"id": 1, "file_name": "image.jpg", "width": 10, "height": 10}
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [1, 2, 3, 4],
"iscrowd": 1,
"area": 12,
}
],
}
),
encoding="utf-8",
)
_, image_paths, annotations = load_coco_annotations(
images_directory_path=str(images_dir),
annotations_path=str(annotations_path),
use_iscrowd=False,
)
assert image_paths == [str(image_path)]
assert "iscrowd" not in annotations[str(image_path)].data
assert "area" not in annotations[str(image_path)].data
def test_detection_dataset_from_coco_accepts_use_iscrowd_false(tmp_path: Path) -> None:
"""DetectionDataset.from_coco forwards use_iscrowd to the COCO loader."""
images_dir = tmp_path / "images"
images_dir.mkdir()
image_path = images_dir / "image.jpg"
cv2.imwrite(str(image_path), np.zeros((10, 10, 3), dtype=np.uint8))
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(
json.dumps(
{
"categories": [{"id": 1, "name": "object"}],
"images": [
{"id": 1, "file_name": "image.jpg", "width": 10, "height": 10}
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [1, 2, 3, 4],
"iscrowd": 1,
"area": 12,
}
],
}
),
encoding="utf-8",
)
dataset = DetectionDataset.from_coco(
images_directory_path=str(images_dir),
annotations_path=str(annotations_path),
use_iscrowd=False,
)
assert "iscrowd" not in dataset.annotations[str(image_path)].data
assert "area" not in dataset.annotations[str(image_path)].data
def test_detection_dataset_from_coco_preserves_show_progress_positional_arg(
tmp_path: Path,
) -> None:
"""The fourth positional from_coco argument remains show_progress."""
images_dir = tmp_path / "images"
images_dir.mkdir()
image_path = images_dir / "image.jpg"
cv2.imwrite(str(image_path), np.zeros((10, 10, 3), dtype=np.uint8))
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(
json.dumps(
{
"categories": [{"id": 1, "name": "object"}],
"images": [
{"id": 1, "file_name": "image.jpg", "width": 10, "height": 10}
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [1, 2, 3, 4],
"iscrowd": 1,
"area": 12,
}
],
}
),
encoding="utf-8",
)
dataset = DetectionDataset.from_coco(
str(images_dir),
str(annotations_path),
False,
False,
)
assert dataset.annotations[str(image_path)].data["iscrowd"].tolist() == [1]
assert dataset.annotations[str(image_path)].data["area"].tolist() == [12.0]
@pytest.mark.parametrize(
(
"image_annotations",
"resolution_wh",
"with_masks",
"use_iscrowd",
"expected_result",
"exception",
),
[
(
[],
(1000, 1000),
False,
False,
Detections.empty(),
DoesNotRaise(),
), # empty image annotations
(
[],
(1000, 1000),
False,
True,
Detections.empty(),
DoesNotRaise(),
), # empty image annotations
(
[
mock_coco_annotation(
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
)
],
(1000, 1000),
False,
False,
Detections(
xyxy=np.array([[0, 0, 100, 100]], dtype=np.float32),
class_id=np.array([0], dtype=int),
data={"coco_raw_segmentation": _empty_raw_segs(1)},
),
DoesNotRaise(),
), # single image annotations
(
[
mock_coco_annotation(
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
)
],
(1000, 1000),
False,
True,
Detections(
xyxy=np.array([[0, 0, 100, 100]], dtype=np.float32),
class_id=np.array([0], dtype=int),
data={
"iscrowd": np.array([0], dtype=int),
"area": np.array([100 * 100]),
"coco_raw_segmentation": _empty_raw_segs(1),
},
),
DoesNotRaise(),
),
(
[
mock_coco_annotation(
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
),
mock_coco_annotation(
category_id=0, bbox=(100, 100, 100, 100), area=100 * 100
),
],
(1000, 1000),
False,
False,
Detections(
xyxy=np.array(
[[0, 0, 100, 100], [100, 100, 200, 200]], dtype=np.float32
),
class_id=np.array([0, 0], dtype=int),
data={"coco_raw_segmentation": _empty_raw_segs(2)},
),
DoesNotRaise(),
), # two image annotations
(
[
mock_coco_annotation(
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
),
mock_coco_annotation(
category_id=0, bbox=(100, 100, 100, 100), area=100 * 100
),
],
(1000, 1000),
False,
True,
Detections(
xyxy=np.array(
[[0, 0, 100, 100], [100, 100, 200, 200]], dtype=np.float32
),
class_id=np.array([0, 0], dtype=int),
data={
"iscrowd": np.array([0, 0], dtype=int),
"area": np.array([100 * 100, 100 * 100]),
"coco_raw_segmentation": _empty_raw_segs(2),
},
),
DoesNotRaise(),
),
(
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 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, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
]
],
dtype=bool,
),
),
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(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
)
],
(5, 5),
True,
True,
Detections(
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array(
[
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
]
],
dtype=bool,
),
data={"iscrowd": np.array([0], dtype=int), "area": np.array([25])},
),
DoesNotRaise(),
),
(
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation={
"size": [5, 5],
"counts": [0, 15, 2, 3, 2, 3],
},
iscrowd=True,
)
],
(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, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
]
],
dtype=bool,
),
),
DoesNotRaise(),
), # single image annotations with mask, RLE segmentation mask
(
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation={
"size": [5, 5],
"counts": [0, 15, 2, 3, 2, 3],
},
iscrowd=True,
)
],
(5, 5),
True,
True,
Detections(
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array(
[
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
]
],
dtype=bool,
),
data={"iscrowd": np.array([1], dtype=int), "area": np.array([25])},
),
DoesNotRaise(),
),
(
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
),
mock_coco_annotation(
category_id=0,
bbox=(3, 0, 2, 2),
area=2 * 2,
segmentation={
"size": [5, 5],
"counts": [15, 2, 3, 2, 3],
},
iscrowd=True,
),
],
(5, 5),
True,
False,
Detections(
xyxy=np.array([[0, 0, 5, 5], [3, 0, 5, 2]], dtype=np.float32),
class_id=np.array([0, 0], dtype=int),
mask=np.array(
[
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
],
[
[0, 0, 0, 1, 1],
[0, 0, 0, 1, 1],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
],
],
dtype=bool,
),
),
DoesNotRaise(),
), # two image annotations with mask, one mask as polygon and second as RLE
(
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
),
mock_coco_annotation(
category_id=0,
bbox=(3, 0, 2, 2),
area=2 * 2,
segmentation={
"size": [5, 5],
"counts": [15, 2, 3, 2, 3],
},
iscrowd=True,
),
],
(5, 5),
True,
True,
Detections(
xyxy=np.array([[0, 0, 5, 5], [3, 0, 5, 2]], dtype=np.float32),
class_id=np.array([0, 0], dtype=int),
mask=np.array(
[
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
],
[
[0, 0, 0, 1, 1],
[0, 0, 0, 1, 1],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
],
],
dtype=bool,
),
data={
"iscrowd": np.array([0, 1], dtype=int),
"area": np.array([25, 4]),
},
),
DoesNotRaise(),
), # two image annotations with mask, one mask as polygon with iscrowd,
# and second as RLE without iscrowd
(
[
mock_coco_annotation(
category_id=0,
bbox=(3, 0, 2, 2),
area=2 * 2,
segmentation={
"size": [5, 5],
"counts": [15, 2, 3, 2, 3],
},
iscrowd=True,
),
mock_coco_annotation(
category_id=1,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
),
],
(5, 5),
True,
False,
Detections(
xyxy=np.array([[3, 0, 5, 2], [0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([0, 1], dtype=int),
mask=np.array(
[
[
[0, 0, 0, 1, 1],
[0, 0, 0, 1, 1],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
],
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
],
],
dtype=bool,
),
),
DoesNotRaise(),
), # two image annotations with mask, first mask as RLE and second as polygon
(
[
mock_coco_annotation(
category_id=0,
bbox=(3, 0, 2, 2),
area=2 * 2,
segmentation={
"size": [5, 5],
"counts": [15, 2, 3, 2, 3],
},
iscrowd=True,
),
mock_coco_annotation(
category_id=1,
bbox=(0, 0, 5, 5),
area=5 * 5,
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
),
],
(5, 5),
True,
True,
Detections(
xyxy=np.array([[3, 0, 5, 2], [0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([0, 1], dtype=int),
mask=np.array(
[
[
[0, 0, 0, 1, 1],
[0, 0, 0, 1, 1],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
],
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
[1, 1, 1, 1, 1],
],
],
dtype=bool,
),
data={
"iscrowd": np.array([1, 0], dtype=int),
"area": np.array([4, 25]),
},
),
DoesNotRaise(),
), # two image annotations with mask, first mask as RLE with is crowd,
# and second as polygon without iscrowd
(
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 4, 4),
area=4 * 4,
segmentation={
"size": [4, 4],
"counts": "52203",
},
iscrowd=True,
)
],
(4, 4),
True,
False,
Detections(
xyxy=np.array([[0, 0, 4, 4]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array(
[
[
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
]
],
dtype=bool,
),
),
DoesNotRaise(),
), # single iscrowd annotation with compressed COCO RLE string counts
],
)
def test_coco_annotations_to_detections(
image_annotations: list[dict],
resolution_wh: tuple[int, int],
with_masks: bool,
use_iscrowd: 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,
use_iscrowd=use_iscrowd,
)
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
@pytest.mark.parametrize(
("detections", "image_id", "annotation_id", "expected_result", "exception"),
[
(
Detections(
xyxy=np.array([[0, 0, 100, 100]], dtype=np.float32),
class_id=np.array([0], dtype=int),
),
0,
0,
[
mock_coco_annotation(
category_id=1, bbox=(0, 0, 100, 100), area=100 * 100
)
],
DoesNotRaise(),
), # no segmentation mask; internal class_id 0 -> COCO category_id 1
(
Detections(
xyxy=np.array([[0, 0, 4, 5]], 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,
),
),
0,
0,
[
mock_coco_annotation(
category_id=1,
bbox=(0, 0, 4, 5),
area=4 * 5,
segmentation=[[0, 0, 0, 4, 3, 4, 3, 0]],
)
],
DoesNotRaise(),
), # segmentation mask in single component,no holes in mask,
# expects polygon mask
(
Detections(
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array(
[
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[0, 0, 0, 1, 1],
[0, 0, 0, 1, 1],
]
],
dtype=bool,
),
),
0,
0,
[
mock_coco_annotation(
category_id=1,
bbox=(0, 0, 5, 5),
area=13,
segmentation={
"size": [5, 5],
"counts": [0, 3, 2, 3, 2, 3, 5, 2, 3, 2],
},
iscrowd=True,
)
],
DoesNotRaise(),
), # segmentation mask with 2 components, no holes in mask, expects RLE mask
(
Detections(
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array(
[
[
[0, 1, 1, 1, 1],
[0, 1, 1, 1, 1],
[1, 1, 0, 0, 1],
[1, 1, 0, 0, 1],
[1, 1, 1, 1, 1],
]
],
dtype=bool,
),
),
0,
0,
[
mock_coco_annotation(
category_id=1,
bbox=(0, 0, 5, 5),
area=19,
segmentation={
"size": [5, 5],
"counts": [2, 10, 2, 3, 2, 6],
},
iscrowd=True,
)
],
DoesNotRaise(),
), # seg mask in single component, with holes in mask, expects RLE mask
],
)
def test_detections_to_coco_annotations(
detections: Detections,
image_id: int,
annotation_id: int,
expected_result: list[dict],
exception: Exception,
) -> None:
with exception:
result, _ = detections_to_coco_annotations(
detections=detections,
image_id=image_id,
annotation_id=annotation_id,
)
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
_DISJOINT_2X2_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,
)
_SINGLE_COMPONENT_MASK = np.array(
[
[
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
]
],
dtype=bool,
)
def _make_iscrowd0_detections(mask: np.ndarray) -> Detections:
"""Build a single-detection Detections with iscrowd=0 from a (1, H, W) mask."""
_, h, w = mask.shape
return Detections(
xyxy=np.array([[0, 0, w, h]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=mask,
data={"iscrowd": np.array([0])},
)
@pytest.mark.parametrize(
("mask", "expected_segment_count"),
[
pytest.param(_DISJOINT_2X2_MASK, 2, id="disjoint-two-parts"),
pytest.param(_SINGLE_COMPONENT_MASK, 1, id="single-component"),
],
)
def test_detections_to_coco_annotations_segmentation_count(
mask: np.ndarray, expected_segment_count: int
) -> None:
"""Non-crowd mask export produces one polygon list entry per connected component."""
annotations, _ = detections_to_coco_annotations(
detections=_make_iscrowd0_detections(mask), image_id=0, annotation_id=0
)
segmentation = annotations[0]["segmentation"]
assert annotations[0]["iscrowd"] == 0
assert isinstance(segmentation, list)
assert len(segmentation) == expected_segment_count
assert all(len(part) >= 6 for part in segmentation)
assert all(np.isfinite(c) for part in segmentation for c in part)
def test_detections_to_coco_annotations_round_trip_disjoint_mask() -> None:
"""Two-part disjoint mask round-trips through COCO export and import unchanged."""
W, H = 5, 5
annotations, _ = detections_to_coco_annotations(
detections=_make_iscrowd0_detections(_DISJOINT_2X2_MASK),
image_id=0,
annotation_id=0,
)
reloaded = coco_annotations_to_masks([annotations[0]], resolution_wh=(W, H))
assert np.array_equal(reloaded[0], _DISJOINT_2X2_MASK[0])
def test_detections_to_coco_annotations_preserves_area_from_data() -> None:
"""area stored in detections.data should be used instead of bbox area."""
detections = Detections(
xyxy=np.array([[10.0, 20.0, 110.0, 120.0]], dtype=np.float32),
class_id=np.array([0], dtype=int),
data={"iscrowd": np.array([0], dtype=int), "area": np.array([5000.0])},
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert len(annotations) == 1
assert annotations[0]["area"] == 5000.0
assert annotations[0]["iscrowd"] == 0
assert type(annotations[0]["iscrowd"]) is int
def test_detections_to_coco_annotations_preserves_iscrowd_from_data_when_no_mask() -> (
None
):
"""iscrowd stored in detections.data should be used when no mask is present."""
detections = Detections(
xyxy=np.array([[0.0, 0.0, 100.0, 100.0]], dtype=np.float32),
class_id=np.array([0], dtype=int),
data={"iscrowd": np.array([1], dtype=int), "area": np.array([1234.5])},
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert len(annotations) == 1
assert annotations[0]["iscrowd"] == 1
assert type(annotations[0]["iscrowd"]) is int
assert annotations[0]["area"] == 1234.5
def test_detections_to_coco_annotations_iscrowd_is_int_when_mask_provided() -> None:
"""iscrowd should be stored as int (0 or 1), not as Python bool."""
mask = np.zeros((1, 5, 5), dtype=bool)
mask[0, 0:3, 0:3] = True # simple single-component rectangle
detections = Detections(
xyxy=np.array([[0.0, 0.0, 3.0, 3.0]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=mask,
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert len(annotations) == 1
assert annotations[0]["iscrowd"] == 0
assert type(annotations[0]["iscrowd"]) is int
def test_detections_to_coco_annotations_data_area_overrides_bbox_with_mask() -> None:
"""data["area"] should override computed bbox area even when a mask is present."""
mask = np.zeros((1, 10, 10), dtype=bool)
mask[0, 0:4, 0:4] = True # 16-pixel polygon area
detections = Detections(
xyxy=np.array([[0.0, 0.0, 10.0, 10.0]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=mask,
data={"area": np.array([999.0])},
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert len(annotations) == 1
assert annotations[0]["area"] == 999.0
def test_detections_to_coco_annotations_mask_area_when_no_data() -> None:
"""Masked detections without stored area export pixel area, not bbox area."""
mask = np.zeros((1, 10, 10), dtype=bool)
mask[0, 0:4, 0:4] = True
detections = Detections(
xyxy=np.array([[0.0, 0.0, 10.0, 10.0]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=mask,
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert len(annotations) == 1
assert annotations[0]["area"] == 16.0
def test_missing_coco_area_with_mask_exports_mask_pixel_area() -> None:
"""Segmented COCO annotations without area export decoded mask area."""
detections = coco_annotations_to_detections(
image_annotations=[
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [0, 0, 10, 10],
"segmentation": [[0, 0, 5, 0, 5, 5, 0, 5]],
}
],
resolution_wh=(10, 10),
with_masks=True,
use_iscrowd=True,
)
assert detections.mask is not None
expected_area = float(np.count_nonzero(detections.mask[0]))
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert annotations[0]["area"] == expected_area
assert annotations[0]["area"] != 100.0
def test_detections_to_coco_annotations_fallback_area_when_no_data() -> None:
"""Box-only detections with no area in data fall back to bbox area."""
detections = Detections(
xyxy=np.array([[10.0, 20.0, 110.0, 120.0]], dtype=np.float32),
class_id=np.array([0], dtype=int),
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert len(annotations) == 1
assert annotations[0]["area"] == 100.0 * 100.0
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:
images_directory = tmp_path / "images"
images_directory.mkdir()
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(
json.dumps(coco_data_with_and_without_segmentation), encoding="utf-8"
)
classes, images, annotations = load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
force_masks=False,
use_iscrowd=True,
)
assert classes == ["object"]
assert len(images) == 4
with_segmentation_path = str(images_directory / "with_segmentation.jpg")
with_segmentation = annotations[with_segmentation_path]
assert with_segmentation.mask is not None
assert with_segmentation.mask.shape == (2, 5, 5)
assert np.array_equal(with_segmentation.data["iscrowd"], np.array([0, 1]))
with_polygon_segmentation_path = str(
images_directory / "with_polygon_segmentation.jpg"
)
with_polygon_segmentation = annotations[with_polygon_segmentation_path]
assert with_polygon_segmentation.mask is not None
assert with_polygon_segmentation.mask.shape == (1, 5, 5)
assert with_polygon_segmentation.mask[0].any()
without_segmentation_path = str(images_directory / "without_segmentation.jpg")
without_segmentation = annotations[without_segmentation_path]
assert without_segmentation.mask is None
assert np.array_equal(
without_segmentation.xyxy, np.array([[0, 0, 2, 2]], dtype=np.float32)
)
without_annotations_path = str(images_directory / "without_annotations.jpg")
assert annotations[without_annotations_path] == Detections.empty()
def test_load_coco_annotations_force_masks_with_no_annotations(
tmp_path, coco_data_with_unannotated_image: dict[str, object]
) -> None:
images_directory = tmp_path / "images"
images_directory.mkdir()
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(
json.dumps(coco_data_with_unannotated_image),
encoding="utf-8",
)
_, _, annotations = load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
force_masks=True,
)
has_segmentation_path = str(images_directory / "has_segmentation.jpg")
has_segmentation = annotations[has_segmentation_path]
assert has_segmentation.mask is not None
assert has_segmentation.mask.shape == (1, 5, 5)
no_annotations_path = str(images_directory / "no_annotations.jpg")
assert annotations[no_annotations_path] == Detections.empty()
def test_coco_annotations_to_masks_handles_rle_polygon_and_invalid_dict() -> None:
image_annotations = [
{"id": 1, "segmentation": {"size": [5, 5], "counts": [15, 2, 3, 2, 3]}},
{"id": 2, "segmentation": [[0, 0, 4, 0, 4, 4, 0, 4]]},
{"id": 3, "segmentation": {"size": [5, 5]}},
]
with pytest.warns(
UserWarning,
match=(
"Skipping annotation 3: segmentation is a dict but missing 'counts' key "
r"\(expected RLE format\)"
),
):
masks = coco_annotations_to_masks(
image_annotations=image_annotations,
resolution_wh=(5, 5),
)
assert masks.shape == (3, 5, 5)
assert masks[0].any()
assert masks[1].any()
assert not masks[2].any()
@pytest.mark.parametrize(
"file_name",
[".", "", "subdir/.."],
)
def test_load_coco_annotations_rejects_file_name_resolving_to_images_directory(
tmp_path,
file_name: str,
) -> None:
"""Reject file_name resolving to the images directory itself (equality guard)."""
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": file_name, "width": 5, "height": 5}],
"annotations": [],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
with pytest.raises(ValueError, match="resolves to the images directory itself"):
load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
@pytest.mark.parametrize(
"malicious_file_name",
[
"../escape.txt",
"../../escape.txt",
"subdir/../../escape.txt",
],
)
def test_load_coco_annotations_rejects_file_name_outside_images_directory(
tmp_path,
malicious_file_name: str,
) -> None:
"""Reject relative traversal file_name values that escape the images directory."""
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": malicious_file_name,
"width": 5,
"height": 5,
}
],
"annotations": [],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
with pytest.raises(ValueError, match="outside the images directory"):
load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
def test_load_coco_annotations_rejects_absolute_file_name(tmp_path) -> None:
"""Reject absolute file_name values that escape the images directory."""
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": "/etc/passwd",
"width": 5,
"height": 5,
}
],
"annotations": [],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
with pytest.raises(ValueError, match="outside the images directory"):
load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
def test_load_coco_annotations_rejects_file_name_resolving_to_directory(
tmp_path,
) -> None:
"""Reject file_name resolving to a subdirectory inside images/ (is_dir guard)."""
images_directory = tmp_path / "images"
images_directory.mkdir()
(images_directory / "subdir").mkdir()
annotations_path = tmp_path / "annotations.json"
coco_data = {
"categories": [{"id": 1, "name": "object", "supercategory": "none"}],
"images": [{"id": 1, "file_name": "subdir", "width": 5, "height": 5}],
"annotations": [],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
with pytest.raises(ValueError, match="resolves to directory"):
load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
def test_load_coco_annotations_rejects_unresolvable_file_name(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Reject file_name values whose resolved path cannot be computed."""
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": "bad.jpg", "width": 5, "height": 5}],
"annotations": [],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
original_resolve = Path.resolve
def fake_resolve(self: Path, *args: object, **kwargs: object) -> Path:
if self == images_directory / "bad.jpg":
raise OSError("unresolvable path")
return original_resolve(self, *args, **kwargs)
monkeypatch.setattr(Path, "resolve", fake_resolve)
with pytest.raises(ValueError, match="invalid path"):
load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
def test_load_coco_annotations_accepts_valid_nested_file_name(tmp_path) -> None:
"""Accept a legitimate nested file_name inside images/ without raising."""
images_directory = tmp_path / "images"
images_directory.mkdir()
(images_directory / "train").mkdir()
annotations_path = tmp_path / "annotations.json"
coco_data = {
"categories": [{"id": 1, "name": "object", "supercategory": "none"}],
"images": [{"id": 1, "file_name": "train/image.jpg", "width": 5, "height": 5}],
"annotations": [],
}
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),
)
expected_path = str(images_directory / "train" / "image.jpg")
assert expected_path in annotations
def test_load_coco_annotations_rejects_duplicate_resolved_file_names(
tmp_path: Path,
) -> None:
"""Aliases for the same file resolve to one canonical COCO entry."""
images_directory = tmp_path / "images"
images_directory.mkdir()
(images_directory / "nested").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},
{"id": 2, "file_name": "nested/../image.jpg", "width": 5, "height": 5},
],
"annotations": [],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
with pytest.raises(ValueError, match="duplicate entries for image"):
load_coco_annotations(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
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))
@pytest.fixture
def coco_data_with_multi_segment_segmentation() -> dict[str, object]:
return {
"categories": [
{
"id": 1,
"name": "cat_eye",
"supercategory": "animal_parts",
}
],
"images": [
{
"id": 1,
"file_name": "image.jpg",
"width": 5,
"height": 5,
}
],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
# bbox spans both segments; area = sum of two 1x1 polygon areas
"bbox": [0, 0, 5, 5],
"area": 2,
"segmentation": [
[0, 0, 1, 0, 1, 1, 0, 1],
[3, 3, 4, 3, 4, 4, 3, 4],
],
"iscrowd": 0,
}
],
}
class TestFromCocoMasks:
"""Integration: DetectionDataset.from_coco loads multi-segment masks."""
@pytest.mark.parametrize("force_masks", [False, True])
def test_multi_segment_masks_merged(
self,
tmp_path,
coco_data_with_multi_segment_segmentation: dict[str, object],
force_masks: bool,
) -> None:
"""Multi-segment masks merge correctly for both force_masks values."""
images_directory = tmp_path / "images"
images_directory.mkdir()
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(
json.dumps(coco_data_with_multi_segment_segmentation), encoding="utf-8"
)
dataset = DetectionDataset.from_coco(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
force_masks=force_masks,
)
annotation = dataset.annotations[str(images_directory / "image.jpg")]
assert annotation.mask is not None
assert annotation.mask.shape == (1, 5, 5)
np.testing.assert_array_equal(
annotation.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,
),
)
def test_multi_segment_masks_uneven_length_no_value_error(self, tmp_path) -> None:
"""Uneven-length segments load without ValueError (issue #1209 regression)."""
images_directory = tmp_path / "images"
images_directory.mkdir()
annotations_path = tmp_path / "annotations.json"
coco_data = {
"categories": [
{"id": 1, "name": "cat_eye", "supercategory": "animal_parts"}
],
"images": [{"id": 1, "file_name": "image.jpg", "width": 5, "height": 5}],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [0, 0, 5, 5],
"area": 2,
"segmentation": [
[0, 0, 1, 0, 1, 1, 0, 1], # 4 points (8 coords)
[3, 3, 4, 3, 4, 4, 3, 4, 2, 4], # 5 points (10 coords)
],
"iscrowd": 0,
}
],
}
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
dataset = DetectionDataset.from_coco(
images_directory_path=str(images_directory),
annotations_path=str(annotations_path),
)
annotation = dataset.annotations[str(images_directory / "image.jpg")]
assert annotation.mask is not None
assert annotation.mask.shape == (1, 5, 5)
# --- category_id 1-indexing (regression for #1181) ---
@pytest.mark.parametrize(
("classes", "expected_ids"),
[
([], []), # empty classes
(["object"], [1]), # single class starts at 1
(["cat", "dog", "bird"], [1, 2, 3]), # ids are sequential and 1-indexed
],
)
def test_classes_to_coco_categories_ids_start_at_one(
classes: list[str], expected_ids: list[int]
) -> None:
"""COCO categories[].id must be 1-indexed (COCO spec / CVAT requirement)."""
categories = classes_to_coco_categories(classes=classes)
assert [category["id"] for category in categories] == expected_ids
def test_detections_to_coco_annotations_category_id_is_one_indexed() -> None:
"""Internal class_id k must serialize to COCO category_id k + 1."""
detections = Detections(
xyxy=np.array([[0, 0, 10, 10], [5, 5, 15, 15], [1, 1, 4, 4]], dtype=np.float32),
class_id=np.array([0, 1, 2], dtype=int),
)
annotations, _ = detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
assert [annotation["category_id"] for annotation in annotations] == [1, 2, 3]
def test_coco_round_trip_preserves_class_ids_and_writes_one_indexed_categories(
tmp_path,
) -> None:
"""as_coco -> from_coco is lossless for internal class_ids while the
on-disk COCO category ids are 1-indexed (regression for #1181)."""
classes = ["cat", "dog"]
image_paths: list[str] = []
annotations: dict[str, Detections] = {}
expected_class_ids = {}
for index, class_id in enumerate([0, 1]):
path = str(tmp_path / f"image_{index}.jpg")
assert cv2.imwrite(path, np.zeros((10, 10, 3), dtype=np.uint8))
image_paths.append(path)
detections = Detections(
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
class_id=np.array([class_id], dtype=int),
)
annotations[path] = detections
expected_class_ids[Path(path).name] = class_id
dataset = DetectionDataset(
classes=classes, images=image_paths, annotations=annotations
)
annotation_path = tmp_path / "annotations.json"
dataset.as_coco(annotations_path=str(annotation_path))
# On-disk COCO ids are 1-indexed.
with open(annotation_path) as f:
payload = json.load(f)
assert sorted(category["id"] for category in payload["categories"]) == [1, 2]
assert sorted(ann["category_id"] for ann in payload["annotations"]) == [1, 2]
# Reading back preserves internal 0-indexed class_ids losslessly.
loaded = DetectionDataset.from_coco(
images_directory_path=str(tmp_path),
annotations_path=str(annotation_path),
)
assert loaded.classes == classes
for image_path, _, detections in loaded:
name = Path(image_path).name
assert detections.class_id is not None
assert detections.class_id.tolist() == [expected_class_ids[name]]
# --- save_coco_annotations: cross-split id chaining (regression for #768) ---
def _tiny_detection_dataset(
tmp_path, prefix: str, num_images: int, dets_per_image: int
) -> DetectionDataset:
"""Build a DetectionDataset of ``num_images`` 10x10 RGB images on disk,
each holding ``dets_per_image`` 1x1 detections of class 0. Image content
is irrelevant; only the per-image Detections drive the COCO write path."""
classes = ["object"]
image_paths: list[str] = []
annotations: dict[str, Detections] = {}
for i in range(num_images):
path = str(tmp_path / f"{prefix}_{i}.jpg")
assert cv2.imwrite(path, np.zeros((10, 10, 3), dtype=np.uint8))
image_paths.append(path)
xyxy = np.array(
[[float(j), 0.0, float(j) + 1.0, 1.0] for j in range(dets_per_image)],
dtype=float,
).reshape(-1, 4)
annotations[path] = Detections(
xyxy=xyxy,
class_id=np.zeros(dets_per_image, dtype=int),
confidence=np.ones(dets_per_image, dtype=float),
)
return DetectionDataset(
classes=classes, images=image_paths, annotations=annotations
)
def _read_ids(annotation_path) -> tuple[list[int], list[int]]:
with open(annotation_path) as f:
payload = json.load(f)
image_ids = [img["id"] for img in payload["images"]]
annotation_ids = [ann["id"] for ann in payload["annotations"]]
return image_ids, annotation_ids
class TestSaveCocoAnnotationsCollisionGuard:
"""COCO export must reject same-basename images before writing."""
def test_raises_on_duplicate_image_basenames(self, tmp_path: Path) -> None:
"""Duplicate image basenames are rejected instead of being collapsed."""
image_paths = []
annotations: dict[str, Detections] = {}
for parent in ("dir_a", "dir_b"):
image_path = tmp_path / parent / "img.jpg"
image_path.parent.mkdir(parents=True, exist_ok=True)
assert cv2.imwrite(str(image_path), np.zeros((10, 10, 3), dtype=np.uint8))
image_path_str = str(image_path)
image_paths.append(image_path_str)
annotations[image_path_str] = Detections.empty()
dataset = DetectionDataset(
classes=["object"], images=image_paths, annotations=annotations
)
with pytest.raises(ValueError, match="COCO image file"):
save_coco_annotations(
dataset=dataset,
annotation_path=str(tmp_path / "annotations.json"),
)
def test_save_coco_annotations_defaults_start_at_one(tmp_path):
dataset = _tiny_detection_dataset(tmp_path, "img", num_images=2, dets_per_image=3)
annotation_path = tmp_path / "annotations.json"
next_image_id, next_annotation_id = save_coco_annotations(
dataset=dataset, annotation_path=str(annotation_path)
)
image_ids, annotation_ids = _read_ids(annotation_path)
assert image_ids == [1, 2]
assert annotation_ids == [1, 2, 3, 4, 5, 6]
# Returned ids are one greater than the highest written, ready to chain.
assert next_image_id == 3
assert next_annotation_id == 7
def test_save_coco_annotations_respects_starting_ids(tmp_path):
dataset = _tiny_detection_dataset(tmp_path, "img", num_images=2, dets_per_image=2)
annotation_path = tmp_path / "annotations.json"
next_image_id, next_annotation_id = save_coco_annotations(
dataset=dataset,
annotation_path=str(annotation_path),
starting_image_id=100,
starting_annotation_id=500,
)
image_ids, annotation_ids = _read_ids(annotation_path)
assert image_ids == [100, 101]
assert annotation_ids == [500, 501, 502, 503]
assert next_image_id == 102
assert next_annotation_id == 504
def test_as_coco_chains_ids_across_splits_without_collision(tmp_path):
"""Regression for #768: exporting train/valid/test splits with the
returned ids fed forward yields globally unique image and annotation ids."""
train = _tiny_detection_dataset(tmp_path, "train", num_images=3, dets_per_image=2)
valid = _tiny_detection_dataset(tmp_path, "valid", num_images=2, dets_per_image=4)
test = _tiny_detection_dataset(tmp_path, "test", num_images=1, dets_per_image=5)
train_path = tmp_path / "train.json"
valid_path = tmp_path / "valid.json"
test_path = tmp_path / "test.json"
next_image_id, next_annotation_id = train.as_coco(annotations_path=str(train_path))
next_image_id, next_annotation_id = valid.as_coco(
annotations_path=str(valid_path),
starting_image_id=next_image_id,
starting_annotation_id=next_annotation_id,
)
test.as_coco(
annotations_path=str(test_path),
starting_image_id=next_image_id,
starting_annotation_id=next_annotation_id,
)
all_image_ids: list[int] = []
all_annotation_ids: list[int] = []
for path in (train_path, valid_path, test_path):
image_ids, annotation_ids = _read_ids(path)
all_image_ids.extend(image_ids)
all_annotation_ids.extend(annotation_ids)
assert len(all_image_ids) == len(set(all_image_ids)), (
"image ids collide across splits"
)
assert len(all_annotation_ids) == len(set(all_annotation_ids)), (
"annotation ids collide across splits"
)
# Concrete chained values.
assert all_image_ids == [1, 2, 3, 4, 5, 6]
assert all_annotation_ids == list(range(1, 6 + 8 + 5 + 1))
def test_save_coco_annotations_empty_dataset_returns_starting_ids(tmp_path):
"""An empty dataset writes a valid (but empty) COCO file and returns
the starting ids unchanged so chaining still composes around it."""
dataset = DetectionDataset(classes=["object"], images=[], annotations={})
annotation_path = tmp_path / "annotations.json"
next_image_id, next_annotation_id = save_coco_annotations(
dataset=dataset,
annotation_path=str(annotation_path),
starting_image_id=7,
starting_annotation_id=42,
)
image_ids, annotation_ids = _read_ids(annotation_path)
assert image_ids == []
assert annotation_ids == []
assert next_image_id == 7
assert next_annotation_id == 42
def test_as_coco_without_annotations_path_returns_starting_ids(tmp_path):
"""When only writing images, the starting ids round-trip unchanged so
chaining still works in the images-only branch."""
dataset = _tiny_detection_dataset(tmp_path, "img", num_images=2, dets_per_image=1)
next_image_id, next_annotation_id = dataset.as_coco(
images_directory_path=str(tmp_path / "imgs"),
starting_image_id=42,
starting_annotation_id=99,
)
assert next_image_id == 42
assert next_annotation_id == 99
def test_save_coco_annotations_annotation_image_id_references_correct_image(tmp_path):
"""Every annotation's image_id must reference an image id present in the
same file, even when a non-default starting_image_id is used."""
dataset = _tiny_detection_dataset(tmp_path, "img", num_images=3, dets_per_image=2)
annotation_path = tmp_path / "annotations.json"
save_coco_annotations(
dataset=dataset,
annotation_path=str(annotation_path),
starting_image_id=100,
starting_annotation_id=500,
)
with open(annotation_path) as f:
coco = json.load(f)
image_id_set = {img["id"] for img in coco["images"]}
annotation_image_ids = {ann["image_id"] for ann in coco["annotations"]}
assert annotation_image_ids <= image_id_set, (
"annotation image_id values reference unknown image ids"
)
def test_save_coco_annotations_zero_annotation_images(tmp_path):
"""Dataset with images but zero detections per image: image ids are
assigned sequentially but annotation list stays empty."""
dataset = _tiny_detection_dataset(tmp_path, "img", num_images=2, dets_per_image=0)
annotation_path = tmp_path / "annotations.json"
next_image_id, next_annotation_id = save_coco_annotations(
dataset=dataset, annotation_path=str(annotation_path)
)
image_ids, annotation_ids = _read_ids(annotation_path)
assert image_ids == [1, 2]
assert annotation_ids == []
assert next_image_id == 3
assert next_annotation_id == 1
# --- Regression: legacy 0-indexed COCO files still load correctly (#1181) ---
def test_from_coco_loads_legacy_zero_indexed_category_ids(tmp_path) -> None:
"""COCO files with 0-indexed category ids (written by supervision <=0.28.x)
must still load and produce correct internal 0-indexed class_ids."""
images_dir = tmp_path / "images"
images_dir.mkdir()
img_path = images_dir / "img.jpg"
assert cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8))
coco_data = {
"categories": [
{"id": 0, "name": "cat", "supercategory": "none"},
{"id": 1, "name": "dog", "supercategory": "none"},
],
"images": [{"id": 1, "file_name": "img.jpg", "width": 10, "height": 10}],
"annotations": [
{
"id": 1,
"image_id": 1,
"category_id": 0,
"bbox": [0, 0, 5, 5],
"area": 25,
"iscrowd": 0,
},
{
"id": 2,
"image_id": 1,
"category_id": 1,
"bbox": [1, 1, 3, 3],
"area": 9,
"iscrowd": 0,
},
],
}
annotations_path = tmp_path / "annotations.json"
annotations_path.write_text(json.dumps(coco_data), encoding="utf-8")
dataset = DetectionDataset.from_coco(
images_directory_path=str(images_dir),
annotations_path=str(annotations_path),
)
assert dataset.classes == ["cat", "dog"]
dets = dataset.annotations[str(img_path)]
assert dets.class_id is not None
assert sorted(dets.class_id.tolist()) == [0, 1]
# --- save_coco_annotations ValueError guards ---
@pytest.mark.parametrize(
("starting_image_id", "starting_annotation_id"),
[
(0, 1),
(1, 0),
(0, 0),
],
)
def test_save_coco_annotations_rejects_zero_starting_ids(
tmp_path, starting_image_id: int, starting_annotation_id: int
) -> None:
"""starting_image_id and starting_annotation_id below 1 must raise ValueError."""
dataset = DetectionDataset(classes=["object"], images=[], annotations={})
annotation_path = tmp_path / "annotations.json"
with pytest.raises(ValueError, match="must be >= 1"):
save_coco_annotations(
dataset=dataset,
annotation_path=str(annotation_path),
starting_image_id=starting_image_id,
starting_annotation_id=starting_annotation_id,
)
# --- detections_to_coco_annotations: class_id=None guard ---
def test_detections_to_coco_annotations_raises_when_class_id_is_none() -> None:
"""Detections with no class_id must raise ValueError before +1 arithmetic."""
detections = Detections(
xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32),
class_id=None,
)
with pytest.raises(ValueError, match="class_id"):
detections_to_coco_annotations(
detections=detections,
image_id=1,
annotation_id=1,
)
# --- Round-trip: multi-class-per-image case ---
def test_coco_round_trip_multi_class_single_image(tmp_path) -> None:
"""Single image with two detections of different classes round-trips losslessly."""
img_path = str(tmp_path / "img.jpg")
assert cv2.imwrite(img_path, np.zeros((10, 10, 3), dtype=np.uint8))
dataset = DetectionDataset(
classes=["cat", "dog"],
images=[img_path],
annotations={
img_path: Detections(
xyxy=np.array([[0, 0, 5, 5], [1, 1, 4, 4]], dtype=np.float32),
class_id=np.array([0, 1], dtype=int),
)
},
)
annotation_path = tmp_path / "annotations.json"
dataset.as_coco(annotations_path=str(annotation_path))
with open(annotation_path) as f:
payload = json.load(f)
assert sorted(ann["category_id"] for ann in payload["annotations"]) == [1, 2]
loaded = DetectionDataset.from_coco(
images_directory_path=str(tmp_path),
annotations_path=str(annotation_path),
)
dets = loaded.annotations[img_path]
assert dets.class_id is not None
assert sorted(dets.class_id.tolist()) == [0, 1]
# --- Regression: segmentation round-trip (#2285) ---
def _coco_annotation_with_segmentation(
segmentation: list[list[int]],
bbox: tuple[float, float, float, float] = (0, 0, 5, 5),
area: float = 25,
) -> dict:
return mock_coco_annotation(
annotation_id=1,
image_id=1,
category_id=1,
bbox=bbox,
area=area,
segmentation=segmentation,
)
def _single_image_coco_data(annotation: dict) -> dict[str, object]:
return {
"info": {},
"licenses": [],
"categories": [{"id": 1, "name": "cat", "supercategory": ""}],
"images": [{"id": 1, "file_name": "img.jpg", "width": 10, "height": 10}],
"annotations": [annotation],
}
def test_detections_to_coco_annotations_exports_all_polygons() -> None:
"""All polygons from a multi-component mask must be exported, not just the first."""
# Build a mask with two separate rectangles (disjoint components)
mask = np.zeros((20, 20), dtype=bool)
mask[1:4, 1:4] = True # top-left component
mask[14:18, 14:18] = True # bottom-right component
detections = Detections(
xyxy=np.array([[1, 1, 4, 4]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array([mask]),
data={"iscrowd": np.array([0], dtype=int)},
)
annotations, _ = detections_to_coco_annotations(
detections=detections, image_id=1, annotation_id=1
)
assert len(annotations) == 1
seg = annotations[0]["segmentation"]
# Both components must appear as separate polygon entries (list of lists)
assert isinstance(seg, list), "segmentation must be a list of polygons"
assert len(seg) >= 2
@pytest.mark.parametrize(
("segmentation", "bbox", "area", "expected_min_polygon_count"),
[
pytest.param(
[[0, 0, 4, 0, 4, 4, 0, 4]],
(0, 0, 5, 5),
25,
1,
id="single-polygon",
),
pytest.param(
[[0, 0, 4, 0, 4, 4, 0, 4], [6, 6, 9, 6, 9, 9, 6, 9]],
(0, 0, 9, 9),
32,
2,
id="multi-polygon",
),
],
)
def test_coco_polygon_segmentation_survives_roundtrip(
tmp_path,
segmentation: list[list[int]],
bbox: tuple[float, float, float, float],
area: float,
expected_min_polygon_count: int,
) -> None:
"""COCO polygon segmentation survives the load/export sequence.
1. Write source COCO JSON with polygon segmentation.
2. Load it through DetectionDataset.from_coco().
3. Export it back to COCO JSON with as_coco().
4. Assert the exported segmentation keeps the expected polygon component count.
"""
images_dir = tmp_path / "images"
images_dir.mkdir()
img_path = images_dir / "img.jpg"
assert cv2.imwrite(str(img_path), np.zeros((10, 10, 3), dtype=np.uint8))
# 1. Write source COCO JSON with polygon segmentation.
ann_path = tmp_path / "annotations.json"
ann_path.write_text(
json.dumps(
_single_image_coco_data(
_coco_annotation_with_segmentation(
segmentation=segmentation, bbox=bbox, area=area
)
)
),
encoding="utf-8",
)
# 2. Load it through the internal DetectionDataset representation.
ds = DetectionDataset.from_coco(
images_directory_path=str(images_dir),
annotations_path=str(ann_path),
)
# 3. Export it back to COCO JSON.
out_ann_path = tmp_path / "out_annotations.json"
ds.as_coco(annotations_path=str(out_ann_path))
with open(out_ann_path) as f:
out = json.load(f)
# 4. Assert polygon component count survives the load/export sequence.
assert len(out["annotations"]) == 1
seg = out["annotations"][0]["segmentation"]
assert isinstance(seg, list)
assert len(seg) >= expected_min_polygon_count
def test_coco_raw_segmentation_preserved_when_masks_not_decoded() -> None:
"""When masks are NOT decoded (with_masks=False), raw polygon data stored in
data['segmentation'] is used as a lossless fallback so as_coco() still emits
non-empty segmentation."""
image_annotations = [
_coco_annotation_with_segmentation(segmentation=[[0, 0, 4, 0, 4, 4, 0, 4]])
]
# Load WITHOUT mask decoding — mask must be None
detections = coco_annotations_to_detections(
image_annotations=image_annotations,
resolution_wh=(10, 10),
with_masks=False,
)
assert detections.mask is None
# Raw segmentation must be stored in data for fallback
assert "coco_raw_segmentation" in detections.data
# Export must still produce non-empty segmentation via fallback
annotations, _ = detections_to_coco_annotations(
detections=detections, image_id=1, annotation_id=1
)
assert len(annotations) == 1
assert annotations[0]["segmentation"] != []
def test_coco_iscrowd_mask_exports_as_rle() -> None:
"""Multi-segment mask exports segmentation as RLE dict (iscrowd inferred as 1)."""
mask = np.zeros((10, 10), dtype=bool)
mask[1:3, 1:3] = True # top-left component
mask[7:9, 7:9] = True # bottom-right component (two separate regions)
detections = Detections(
xyxy=np.array([[1, 1, 8, 8]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array([mask]),
)
annotations, _ = detections_to_coco_annotations(
detections=detections, image_id=1, annotation_id=1
)
assert len(annotations) == 1
seg = annotations[0]["segmentation"]
assert isinstance(seg, dict), "multi-segment mask must export as RLE dict, not list"
assert "counts" in seg
assert "size" in seg