add concurrent processing for inference slicer

This commit is contained in:
James Gallagher committed 2023-09-08 12:43:20 +01:00
1 parent 3d4326c79b
commit ea0ec29c81
1 file changed
+26 -5
@@ -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],