add concurrent processing for inference slicer
This commit is contained in:
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],
|
||||
|
||||
Reference in new issue
Block a user