From 7f114cba4c13d621ce86177db1fdbbfa2d4575a2 Mon Sep 17 00:00:00 2001
From: magda skoczen <m.sko@o2.pl>
Date: Fri, 3 May 2024 00:24:46 +0200
Subject: [PATCH] RLE decoding

---
 supervision/dataset/formats/coco.py | 33 +++++++++++++++++++++++++----
 supervision/detection/utils.py      | 10 +++++++++
 test/dataset/formats/test_coco.py   |  2 +-
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py
index 4f8679d5..4105f86e 100644
--- a/supervision/dataset/formats/coco.py
+++ b/supervision/dataset/formats/coco.py
@@ -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))
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index 3eeba5b4..2b6f7d63 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -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')
diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py
index 5055a859..7254c9a9 100644
--- a/test/dataset/formats/test_coco.py
+++ b/test/dataset/formats/test_coco.py
@@ -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),