supervision/test/dataset/formats/test_coco.py

788 lines
26 KiB
Python

from __future__ import annotations
from contextlib import ExitStack as DoesNotRaise
import numpy as np
import pytest
from supervision import Detections
from supervision.dataset.formats.coco import (
build_coco_class_index_mapping,
classes_to_coco_categories,
coco_annotations_to_detections,
coco_categories_to_classes,
detections_to_coco_annotations,
group_coco_annotations_by_image_id,
)
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),
}
@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
@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),
),
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]),
},
),
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),
),
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]),
},
),
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],
]
]
),
),
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, 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],
]
]
),
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],
]
]
),
),
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],
]
]
),
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],
],
]
),
),
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],
],
]
),
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],
],
]
),
),
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],
],
]
),
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
],
)
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=0, bbox=(0, 0, 100, 100), area=100 * 100
)
],
DoesNotRaise(),
), # no segmentation mask
(
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],
]
]
),
),
0,
0,
[
mock_coco_annotation(
category_id=0,
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],
]
]
),
),
0,
0,
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
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],
]
]
),
),
0,
0,
[
mock_coco_annotation(
category_id=0,
bbox=(0, 0, 5, 5),
area=5 * 5,
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