diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 930420d4..f285b4ed 100644 --- a/supervision/annotators/core.py +++ b/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( diff --git a/supervision/detection/tools/inference_slicer.py b/supervision/detection/tools/inference_slicer.py index dee01a80..5f6fb391 100644 --- a/supervision/detection/tools/inference_slicer.py +++ b/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],