Merge pull request roboflow#437 from ashishdatta/from_roboflow_tracker_id

SkalskiP · web-flow · commit 8eb827b31d23 · 2023-10-09T19:58:50.000+02:00
Update sv.Detections.from_roboflow to extract tracker_id values from inference response
diff --git a/supervision/detection/core.py b/supervision/detection/core.py
@@ -426,7 +426,7 @@ def from_roboflow(cls, roboflow_result: dict) -> Detections:
             >>> detections = sv.Detections.from_roboflow(roboflow_result)
             ```
         """
-        xyxy, confidence, class_id, masks = process_roboflow_result(
+        xyxy, confidence, class_id, masks, trackers = process_roboflow_result(
             roboflow_result=roboflow_result
         )
 
@@ -438,6 +438,7 @@ def from_roboflow(cls, roboflow_result: dict) -> Detections:
             confidence=confidence,
             class_id=class_id,
             mask=masks,
+            tracker_id=trackers,
         )
 
     @classmethod
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
@@ -333,14 +333,15 @@ 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]]:
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
     if not roboflow_result["predictions"]:
-        return np.empty((0, 4)), np.empty(0), np.empty(0), None
+        return np.empty((0, 4)), np.empty(0), np.empty(0), None, None
 
     xyxy = []
     confidence = []
     class_id = []
     masks = []
+    tracker_ids = []
 
     image_width = int(roboflow_result["image"]["width"])
     image_height = int(roboflow_result["image"]["height"])
@@ -359,6 +360,8 @@ def process_roboflow_result(
             xyxy.append([x_min, y_min, x_max, y_max])
             class_id.append(prediction["class_id"])
             confidence.append(prediction["confidence"])
+            if "tracker_id" in prediction:
+                tracker_ids.append(prediction["tracker_id"])
         elif len(prediction["points"]) >= 3:
             polygon = np.array(
                 [[point["x"], point["y"]] for point in prediction["points"]], dtype=int
@@ -368,13 +371,16 @@ def process_roboflow_result(
             class_id.append(prediction["class_id"])
             confidence.append(prediction["confidence"])
             masks.append(mask)
+            if "tracker_id" in prediction:
+                tracker_ids.append(prediction["tracker_id"])
 
     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)
     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
 
-    return xyxy, confidence, class_id, masks
+    return xyxy, confidence, class_id, masks, tracker_id
 
 
 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
@@ -264,7 +264,7 @@ def test_filter_polygons_by_area(
     [
         (
             {"predictions": [], "image": {"width": 1000, "height": 1000}},
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None),
+            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
             DoesNotRaise(),
         ),  # empty result
         (
@@ -287,6 +287,7 @@ def test_filter_polygons_by_area(
                 np.array([0.9]),
                 np.array([0]),
                 None,
+                None,
             ),
             DoesNotRaise(),
         ),  # single correct object detection result
@@ -301,6 +302,7 @@ def test_filter_polygons_by_area(
                         "confidence": 0.9,
                         "class_id": 0,
                         "class": "person",
+                        "tracker_id": 1,
                     },
                     {
                         "x": 500.0,
@@ -310,6 +312,7 @@ def test_filter_polygons_by_area(
                         "confidence": 0.8,
                         "class_id": 7,
                         "class": "truck",
+                        "tracker_id": 2,
                     },
                 ],
                 "image": {"width": 1000, "height": 1000},
@@ -319,6 +322,7 @@ def test_filter_polygons_by_area(
                 np.array([0.9, 0.8]),
                 np.array([0, 7]),
                 None,
+                np.array([1, 2]),
             ),
             DoesNotRaise(),
         ),  # two correct object detection result
@@ -334,11 +338,12 @@ def test_filter_polygons_by_area(
                         "class_id": 0,
                         "class": "person",
                         "points": [],
+                        "tracker_id": None,
                     }
                 ],
                 "image": {"width": 1000, "height": 1000},
             },
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None),
+            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no points
         (
@@ -357,7 +362,7 @@ def test_filter_polygons_by_area(
                 ],
                 "image": {"width": 1000, "height": 1000},
             },
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None),
+            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no enough points
         (
@@ -386,6 +391,7 @@ def test_filter_polygons_by_area(
                 np.array([0.9]),
                 np.array([0]),
                 TEST_MASK,
+                None,
             ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no enough points
@@ -425,14 +431,17 @@ def test_filter_polygons_by_area(
                 np.array([0.9]),
                 np.array([0]),
                 TEST_MASK,
+                None,
             ),
             DoesNotRaise(),
         ),  # two instance segmentation results - one correct, one incorrect
     ],
 )
 def test_process_roboflow_result(
     roboflow_result: dict,
-    expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
+    expected_result: Tuple[
+        np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray
+    ],
     exception: Exception,
 ) -> None:
     with exception:
@@ -443,6 +452,9 @@ def test_process_roboflow_result(
         assert (result[3] is None and expected_result[3] is None) or (
             np.array_equal(result[3], expected_result[3])
         )
+        assert (result[4] is None and expected_result[4] is None) or (
+            np.array_equal(result[4], expected_result[4])
+        )
 
 
 @pytest.mark.parametrize(

Original file line number	Diff line number	Diff line change
`@@ -426,7 +426,7 @@ def from_roboflow(cls, roboflow_result: dict) -> Detections:`
`426`	`426`	`>>> detections = sv.Detections.from_roboflow(roboflow_result)`
`427`	`427`	```
`428`	`428`	`"""`
`429`		`- xyxy, confidence, class_id, masks = process_roboflow_result(`
	`429`	`+ xyxy, confidence, class_id, masks, trackers = process_roboflow_result(`
`430`	`430`	`roboflow_result=roboflow_result`
`431`	`431`	`)`
`432`	`432`
`@@ -438,6 +438,7 @@ def from_roboflow(cls, roboflow_result: dict) -> Detections:`
`438`	`438`	`confidence=confidence,`
`439`	`439`	`class_id=class_id,`
`440`	`440`	`mask=masks,`
	`441`	`+ tracker_id=trackers,`
`441`	`442`	`)`
`442`	`443`
`443`	`444`	`@classmethod`