From ea0ec29c8105bb83f4f0fa748a6897c1706f9e50 Mon Sep 17 00:00:00 2001 From: James Gallagher Date: Fri, 8 Sep 2023 12:43:20 +0100 Subject: [PATCH] add concurrent processing for inference slicer --- .../detection/tools/inference_slicer.py | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/supervision/detection/tools/inference_slicer.py b/supervision/detection/tools/inference_slicer.py index 1fbd0697..d598c715 100644 --- a/supervision/detection/tools/inference_slicer.py +++ b/supervision/detection/tools/inference_slicer.py @@ -6,6 +6,8 @@ from supervision.detection.core import Detections, validate_inference_callback from supervision.detection.utils import move_boxes from supervision.utils.image import crop_image +import concurrent.futures + def move_detections(detections: Detections, offset: np.array) -> Detections: """ @@ -97,15 +99,34 @@ 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 concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + futures = [ + executor.submit(self._run_callback, image, offset) for offset in offsets + ] + for future in concurrent.futures.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],