RLE decoding

This commit is contained in:
magda skoczen 2024-05-03 00:24:46 +02:00
parent cc5e72dd5b
commit 7f114cba4c
3 changed files with 40 additions and 5 deletions

View File

@ -11,7 +11,7 @@ from supervision.dataset.utils import (
map_detections_class_id,
)
from supervision.detection.core import Detections
from supervision.detection.utils import polygon_to_mask
from supervision.detection.utils import polygon_to_mask, rle_to_mask
from supervision.utils.file import read_json_file, save_json_file
@ -68,6 +68,26 @@ def _polygons_to_masks(
dtype=bool,
)
def _rles_to_masks(
rles: List[np.ndarray], resolution_wh: Tuple[int, int]
) -> np.ndarray:
return np.array(
[
rle_to_mask(rle=rle, resolution_wh=resolution_wh)
for rle in rles
],
dtype=bool,
)
def _concatenate_annotation_masks(mask_polygon, mask_rle):
if mask_polygon.ndim == 3 and mask_rle.ndim == 3:
return np.concatenate((mask_polygon, mask_rle))
elif mask_polygon.ndim == 3:
return mask_polygon
elif mask_rle.ndim == 3:
return mask_rle
else:
None
def coco_annotations_to_detections(
image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
@ -87,11 +107,16 @@ def coco_annotations_to_detections(
np.reshape(
np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
)
for image_annotation in image_annotations
for image_annotation in image_annotations if not image_annotation["iscrowd"]
]
mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
mask_polygon = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
rles = [np.array(image_annotation["segmentation"]["counts"])
for image_annotation in image_annotations if image_annotation["iscrowd"]]
mask_rle = _rles_to_masks(rles = rles, resolution_wh = resolution_wh)
return Detections(
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=_concatenate_annotation_masks(mask_polygon=mask_polygon, mask_rle=mask_rle)
)
return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))

View File

@ -766,3 +766,13 @@ def get_data_item(
raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
return subset_data
def rle_to_mask(rle: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
width, height = resolution_wh
zero_one_values = np.zeros_like(rle)
zero_one_values[1::2]=1
decoded_rle = np.repeat(zero_one_values, rle)
return decoded_rle.reshape((height,width), order='F')

View File

@ -252,7 +252,7 @@ def test_group_coco_annotations_by_image_id(
[
mock_cock_coco_annotation(
category_id=0, bbox=(0, 0, 10, 10), area=10 * 10,
segmentation = {'size':[20,20], 'counts':[0, 5, 20, 5, 40, 5, 60, 5, 80, 5, 100, 10, 120, 10, 140, 10, 160, 10, 180, 10]}, iscrowd = True
segmentation = {'size':[20,20], 'counts':[0, 10, 10, 10, 10, 10, 10, 10, 10, 10, 15, 5, 15, 5, 15, 5, 15, 5, 15, 5, 210]}, iscrowd = True
)
],
(20, 20),