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

automatic RLE for masks with holes or in multiple pieces

This commit is contained in:
magda skoczen 2024-05-16 08:40:28 +02:00
parent f08404eb7d
commit d8679d9c46
2 changed files with 148 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
supervision/dataset/formats/coco.py
View File

@ -11,6 +11,7 @@ from supervision.dataset.utils import (
approximate_mask_with_polygons,
map_detections_class_id,
rle_to_mask,
mask_to_rle
)
from supervision.detection.core import Detections
from supervision.detection.utils import polygon_to_mask
@ -106,6 +107,21 @@ def coco_annotations_to_detections(
return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
def _mask_has_holes(mask: np.ndarray)-> bool:
_, hierarchy = cv2.findContours(mask.astype(np.uint8), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_SIMPLE)
parent_countour_index = 3
for h in hierarchy[0]:
if h[parent_countour_index] != -1:
return True
return False
def _mask_has_multiple_segments(mask: np.ndarray)-> bool:
number_of_labels, _ = cv2.connectedComponents(mask.astype(np.uint8), connectivity=4)
return number_of_labels > 2
def detections_to_coco_annotations(
detections: Detections,
image_id: int,
@ -118,26 +134,31 @@ def detections_to_coco_annotations(
for xyxy, mask, _, class_id, _, _ in detections:
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
segmentation = []
iscrowd = 0
if mask is not None:
segmentation = list(
approximate_mask_with_polygons(
mask=mask,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)[0].flatten()
)
# todo: flag for when to use RLE?
# segmentation = {"counts": mask_to_rle(binary_mask=mask),
# "size": list(mask.shape[:2])}
iscrowd = _mask_has_holes(mask = mask) or \
_mask_has_multiple_segments(mask = mask)
if iscrowd:
segmentation = {"counts": mask_to_rle(mask=mask),
"size": list(mask.shape[:2])}
else:
segmentation = [list(
approximate_mask_with_polygons(
mask=mask,
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
)[0].flatten()
)] # multicomponent masks supported only for rle format
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": [segmentation] if segmentation else [],
"iscrowd": 0, ## todo: iscrowd depends on flag 1 if RLE 0 if polygon
"segmentation": segmentation,
"iscrowd": iscrowd,
}
coco_annotations.append(coco_annotation)
annotation_id += 1

116
test/dataset/formats/test_coco.py
View File

@ -1,5 +1,5 @@
from contextlib import ExitStack as DoesNotRaise
from typing import Dict, List, Tuple
from typing import Dict, List, Tuple, Union
import numpy as np
import pytest
@ -11,6 +11,7 @@ from supervision.dataset.formats.coco import (
coco_annotations_to_detections,
coco_categories_to_classes,
group_coco_annotations_by_image_id,
detections_to_coco_annotations
)
@ -20,9 +21,11 @@ def mock_cock_coco_annotation(
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] = None,
segmentation: Union[List[list], Dict] = None,
iscrowd: bool = False,
) -> dict:
if not segmentation:
segmentation = []
return {
"id": annotation_id,
"image_id": image_id,
@ -454,3 +457,112 @@ def test_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_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100)],
DoesNotRaise(),
), # no segmentation 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, 1, 1],
# [1, 1, 1, 1, 1],
# [1, 1, 1, 1, 1],
# ]
# ]
# ),
# ),
# 0,
# 0,
# [mock_cock_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]])],
# 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_cock_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_cock_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(),
) # segmentation 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