From e91151618042c2e99c3a91d66cacf9f76dc4470f Mon Sep 17 00:00:00 2001
From: SkalskiP <piotr.skalski92@gmail.com>
Date: Wed, 4 Oct 2023 10:44:49 +0200
Subject: [PATCH] Change prediction class handling from text to ids

Updated the prediction parsing logic in supervision/detection/utils.py and supervision/detection/core.py to use 'class_id' instead of 'class' names. Adapted tests in test/detection/test_utils.py accordingly to reflect the same change.
---
 supervision/detection/core.py  |  1 +
 supervision/detection/utils.py |  4 ++--
 test/detection/test_utils.py   | 13 ++++++++++---
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/supervision/detection/core.py b/supervision/detection/core.py
index 6c6a948e..0ba28491 100644
--- a/supervision/detection/core.py
+++ b/supervision/detection/core.py
@@ -456,6 +456,7 @@ class Detections:
             ...             "y": 0.5,
             ...             "width": 0.2,
             ...             "height": 0.3,
+            ...             "class_id": 0,
             ...             "class": "person",
             ...             "confidence": 0.9
             ...         },
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index 94cbebdd..2beff657 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -357,7 +357,7 @@ def process_roboflow_result(
 
         if "points" not in prediction:
             xyxy.append([x_min, y_min, x_max, y_max])
-            class_id.append(int(prediction["class"]))
+            class_id.append(prediction["class_id"])
             confidence.append(prediction["confidence"])
         elif len(prediction["points"]) >= 3:
             polygon = np.array(
@@ -365,7 +365,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(int(prediction["class"]))
+            class_id.append(prediction["class_id"])
             confidence.append(prediction["confidence"])
             masks.append(mask)
 
diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py
index 330b94f0..529914ae 100644
--- a/test/detection/test_utils.py
+++ b/test/detection/test_utils.py
@@ -260,7 +260,7 @@ def test_filter_polygons_by_area(
 
 
 @pytest.mark.parametrize(
-    "roboflow_result, class_list, expected_result, exception",
+    "roboflow_result, expected_result, exception",
     [
         (
             {"predictions": [], "image": {"width": 1000, "height": 1000}},
@@ -276,6 +276,7 @@ def test_filter_polygons_by_area(
                         "width": 50.0,
                         "height": 50.0,
                         "confidence": 0.9,
+                        "class_id": 0,
                         "class": "person",
                     }
                 ],
@@ -298,6 +299,7 @@ def test_filter_polygons_by_area(
                         "width": 50.0,
                         "height": 50.0,
                         "confidence": 0.9,
+                        "class_id": 0,
                         "class": "person",
                     },
                     {
@@ -306,6 +308,7 @@ def test_filter_polygons_by_area(
                         "width": 100.0,
                         "height": 100.0,
                         "confidence": 0.8,
+                        "class_id": 7,
                         "class": "truck",
                     },
                 ],
@@ -314,7 +317,7 @@ def test_filter_polygons_by_area(
             (
                 np.array([[175.0, 275.0, 225.0, 325.0], [450.0, 450.0, 550.0, 550.0]]),
                 np.array([0.9, 0.8]),
-                np.array([0, 2]),
+                np.array([0, 7]),
                 None,
             ),
             DoesNotRaise(),
@@ -328,6 +331,7 @@ def test_filter_polygons_by_area(
                         "width": 50.0,
                         "height": 50.0,
                         "confidence": 0.9,
+                        "class_id": 0,
                         "class": "person",
                         "points": [],
                     }
@@ -346,6 +350,7 @@ def test_filter_polygons_by_area(
                         "width": 50.0,
                         "height": 50.0,
                         "confidence": 0.9,
+                        "class_id": 0,
                         "class": "person",
                         "points": [{"x": 200.0, "y": 300.0}, {"x": 250.0, "y": 300.0}],
                     }
@@ -364,6 +369,7 @@ def test_filter_polygons_by_area(
                         "width": 50.0,
                         "height": 50.0,
                         "confidence": 0.9,
+                        "class_id": 0,
                         "class": "person",
                         "points": [
                             {"x": 200.0, "y": 300.0},
@@ -392,6 +398,7 @@ def test_filter_polygons_by_area(
                         "width": 50.0,
                         "height": 50.0,
                         "confidence": 0.9,
+                        "class_id": 0,
                         "class": "person",
                         "points": [
                             {"x": 200.0, "y": 300.0},
@@ -406,6 +413,7 @@ def test_filter_polygons_by_area(
                         "width": 100.0,
                         "height": 100.0,
                         "confidence": 0.8,
+                        "class_id": 7,
                         "class": "truck",
                         "points": [],
                     },
@@ -424,7 +432,6 @@ def test_filter_polygons_by_area(
 )
 def test_process_roboflow_result(
     roboflow_result: dict,
-    class_list: List[str],
     expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
     exception: Exception,
 ) -> None: