Merge branch 'roboflow:develop' into new1

supervision/annotators/core.py

@@ -630,6 +630,7 @@ class TraceAnnotator:
         position: Optional[Position] = Position.CENTER,
         trace_length: int = 30,
         thickness: int = 2,
+        color_map: str = "class",
     ):
         """
         Args:
@@ -640,11 +641,14 @@ class TraceAnnotator:
             trace_length (int): The maximum length of the trace in terms of historical
                 points. Defaults to `30`.
             thickness (int): The thickness of the trace lines. Defaults to `2`.
+            color_map (str): Strategy for mapping colors to annotations.
+                Options are `index`, `class`, or `track`.
         """
         self.color: Union[Color, ColorPalette] = color
         self.position = position
         self.trace = Trace(max_size=trace_length)
         self.thickness = thickness
+        self.color_map: ColorMap = ColorMap(color_map)
 
     def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
         """
@@ -677,15 +681,14 @@ class TraceAnnotator:
         """
         self.trace.put(detections)
 
-        for i, (xyxy, mask, confidence, class_id, tracker_id) in enumerate(detections):
-            class_id = detections.class_id[i] if class_id is not None else None
-            idx = class_id if class_id is not None else i
-            color = (
-                self.color.by_idx(idx)
-                if isinstance(self.color, ColorPalette)
-                else self.color
+        for detection_idx in range(len(detections)):
+            tracker_id = int(detections.tracker_id[detection_idx])
+            idx = resolve_color_idx(
+                detections=detections,
+                detection_idx=detection_idx,
+                color_map=self.color_map,
             )
-
+            color = resolve_color(color=self.color, idx=idx)
             xy = self.trace.get(tracker_id=tracker_id)
             if len(xy) > 1:
                 scene = cv2.polylines(

supervision/detection/tools/inference_slicer.py

@@ -1,3 +1,4 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Callable, Optional, Tuple
 
 import numpy as np
@@ -28,15 +29,16 @@ class InferenceSlicer:
     involves dividing a larger image into smaller slices, performing inference on each
     slice, and then merging the detections.
 
-    Attributes:
+    Args:
         slice_wh (Tuple[int, int]): Dimensions of each slice in the format
             `(width, height)`.
         overlap_ratio_wh (Tuple[float, float]): Overlap ratio between consecutive
             slices in the format `(width_ratio, height_ratio)`.
-        iou_threshold (Optional[float]): Intersection over Union (IoU) threshold used
-            for non-max suppression.
-        callback (Callable): A function that performs inference on a given image slice
-            and returns detections.
+        iou_threshold (Optional[float]): Intersection over Union (IoU) threshold
+            used for non-max suppression.
+        callback (Callable): A function that performs inference on a given image
+            slice and returns detections.
+        thread_workers (int): Number of threads for parallel execution.
 
     Note:
         The class ensures that slices do not exceed the boundaries of the original
@@ -51,11 +53,13 @@ class InferenceSlicer:
         slice_wh: Tuple[int, int] = (320, 320),
         overlap_ratio_wh: Tuple[float, float] = (0.2, 0.2),
         iou_threshold: Optional[float] = 0.5,
+        thread_workers: int = 1,
     ):
         self.slice_wh = slice_wh
         self.overlap_ratio_wh = overlap_ratio_wh
         self.iou_threshold = iou_threshold
         self.callback = callback
+        self.thread_workers = thread_workers
         validate_inference_callback(callback=callback)
 
     def __call__(self, image: np.ndarray) -> Detections:
@@ -98,15 +102,35 @@ class InferenceSlicer:
             overlap_ratio_wh=self.overlap_ratio_wh,
         )
 
-        for offset in offsets:
-            image_slice = crop_image(image=image, xyxy=offset)
-            detections = self.callback(image_slice)
-            detections = move_detections(detections=detections, offset=offset[:2])
-            detections_list.append(detections)
+        with ThreadPoolExecutor(max_workers=self.thread_workers) as executor:
+            futures = [
+                executor.submit(self._run_callback, image, offset) for offset in offsets
+            ]
+            for future in as_completed(futures):
+                detections_list.append(future.result())
+
         return Detections.merge(detections_list=detections_list).with_nms(
             threshold=self.iou_threshold
         )
 
+    def _run_callback(self, image, offset) -> Detections:
+        """
+        Run the provided callback on a slice of an image.
+
+        Args:
+            image (np.ndarray): The input image on which inference needs to run
+            offset (np.ndarray): An array of shape `(4,)` containing coordinates
+                for the slice.
+
+        Returns:
+            Detections: A collection of detections for the slice.
+        """
+        image_slice = crop_image(image=image, xyxy=offset)
+        detections = self.callback(image_slice)
+        detections = move_detections(detections=detections, offset=offset[:2])
+
+        return detections
+
     @staticmethod
     def _generate_offset(
         resolution_wh: Tuple[int, int],