RLE decoding
This commit is contained in:
parent
cc5e72dd5b
commit
7f114cba4c
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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')
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Reference in New Issue