From 706ff0fd460cb187cfba2f57f4e0d6017548c4c7 Mon Sep 17 00:00:00 2001
From: SkalskiP <piotr.skalski92@gmail.com>
Date: Tue, 2 Jan 2024 21:00:17 +0100
Subject: [PATCH] initial update of `process_roboflow_result` function

---
 supervision/detection/utils.py | 18 ++++++++++++---
 test/detection/test_utils.py   | 40 +++++++++++++++++++++++++++++++---
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index 2df89a84..3b9d2cb8 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -331,13 +331,21 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
 
 def process_roboflow_result(
     roboflow_result: dict,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
+) -> Tuple[
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    Optional[np.ndarray],
+    np.ndarray,
+    Dict[str, List[np.ndarray]]
+]:
     if not roboflow_result["predictions"]:
-        return np.empty((0, 4)), np.empty(0), np.empty(0), None, None
+        return np.empty((0, 4)), np.empty(0), np.empty(0), None, None, {"class_name": np.empty(0)}
 
     xyxy = []
     confidence = []
     class_id = []
+    class_name = []
     masks = []
     tracker_ids = []
 
@@ -357,6 +365,7 @@ def process_roboflow_result(
         if "points" not in prediction:
             xyxy.append([x_min, y_min, x_max, y_max])
             class_id.append(prediction["class_id"])
+            class_name.append(prediction["class"])
             confidence.append(prediction["confidence"])
             if "tracker_id" in prediction:
                 tracker_ids.append(prediction["tracker_id"])
@@ -367,6 +376,7 @@ def process_roboflow_result(
             mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
             xyxy.append([x_min, y_min, x_max, y_max])
             class_id.append(prediction["class_id"])
+            class_name.append(prediction["class"])
             confidence.append(prediction["confidence"])
             masks.append(mask)
             if "tracker_id" in prediction:
@@ -375,10 +385,12 @@ def process_roboflow_result(
     xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
     confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0)
     class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0)
+    class_name = np.array(class_name) if len(class_name) > 0 else np.empty(0)
     masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
     tracker_id = np.array(tracker_ids).astype(int) if len(tracker_ids) > 0 else None
+    data = {"class_name": class_name}
 
-    return xyxy, confidence, class_id, masks, tracker_id
+    return xyxy, confidence, class_id, masks, tracker_id, data
 
 
 def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py
index a24119ce..423bde54 100644
--- a/test/detection/test_utils.py
+++ b/test/detection/test_utils.py
@@ -268,7 +268,14 @@ def test_filter_polygons_by_area(
     [
         (
             {"predictions": [], "image": {"width": 1000, "height": 1000}},
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None,
+                None,
+                {"class_name": np.empty(0)}
+            ),
             DoesNotRaise(),
         ),  # empty result
         (
@@ -292,6 +299,7 @@ def test_filter_polygons_by_area(
                 np.array([0]),
                 None,
                 None,
+                {"class_name": np.array(["person"])}
             ),
             DoesNotRaise(),
         ),  # single correct object detection result
@@ -327,6 +335,7 @@ def test_filter_polygons_by_area(
                 np.array([0, 7]),
                 None,
                 np.array([1, 2]),
+                {"class_name": np.array(["person", "truck"])}
             ),
             DoesNotRaise(),
         ),  # two correct object detection result
@@ -347,7 +356,14 @@ def test_filter_polygons_by_area(
                 ],
                 "image": {"width": 1000, "height": 1000},
             },
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None,
+                None,
+                {"class_name": np.empty(0)}
+            ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no points
         (
@@ -366,7 +382,14 @@ def test_filter_polygons_by_area(
                 ],
                 "image": {"width": 1000, "height": 1000},
             },
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None,
+                None,
+                {"class_name": np.empty(0)}
+            ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no enough points
         (
@@ -396,6 +419,7 @@ def test_filter_polygons_by_area(
                 np.array([0]),
                 TEST_MASK,
                 None,
+                {"class_name": np.array(["person"])}
             ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no enough points
@@ -436,6 +460,7 @@ def test_filter_polygons_by_area(
                 np.array([0]),
                 TEST_MASK,
                 None,
+                {"class_name": np.array(["person"])}
             ),
             DoesNotRaise(),
         ),  # two instance segmentation results - one correct, one incorrect
@@ -459,6 +484,15 @@ def test_process_roboflow_result(
         assert (result[4] is None and expected_result[4] is None) or (
             np.array_equal(result[4], expected_result[4])
         )
+        for key in result[5]:
+            if isinstance(result[5][key], np.ndarray):
+                assert np.array_equal(
+                    result[5][key], expected_result[5][key]
+                ), f"Mismatch in arrays for key {key}"
+            else:
+                assert (
+                    result[5][key] == expected_result[5][key]
+                ), f"Mismatch in non-array data for key {key}"
 
 
 @pytest.mark.parametrize(