From 9b35a0bda0eb4f1af24d9b02636579d394bb9790 Mon Sep 17 00:00:00 2001 From: ramiro Date: Tue, 28 Jul 2026 01:06:23 -0300 Subject: [PATCH] Feat #576: support NumPy array input in batch_image_detection for MDv5 and RT-DETR batch_image_detection already accepted a list of NumPy arrays in YOLOV8Base, but the other detector backends still required the images to exist on disk, forcing in-memory pipelines (images stored as binary columns in Parquet, PySpark/Dask workflows) to write intermediate JPEG files just to run detection. YOLOV5Base now has the same "handle numpy array input" branch as YOLOV8Base: it batches the arrays, stacks them with the MegaDetector_v5_Transform already used by single_image_detection, and rescales the boxes using each array's own shape. RTDETRApacheBase already iterated one image at a time, so instead of duplicating the loop its image source became a lazy generator: arrays go through Image.fromarray, while the directory case keeps using DetectionImageFolder and opens each file inside the loop exactly as before. In both backends the first parameter was renamed data_path -> data_source to match the YOLOV8Base signature, array inputs get the index as img_id, and there is one result entry per input image. Closes #576 Co-Authored-By: Claude Fable 5 --- .../rtdetr_apache/rtdetr_apache_base.py | 37 +++++++++++-------- .../ultralytics_based/yolov5_base.py | 31 ++++++++++++++-- 2 files changed, 50 insertions(+), 18 deletions(-) diff --git a/PytorchWildlife/models/detection/rtdetr_apache/rtdetr_apache_base.py b/PytorchWildlife/models/detection/rtdetr_apache/rtdetr_apache_base.py index 4e5004e38..ab2f60107 100644 --- a/PytorchWildlife/models/detection/rtdetr_apache/rtdetr_apache_base.py +++ b/PytorchWildlife/models/detection/rtdetr_apache/rtdetr_apache_base.py @@ -5,6 +5,7 @@ # Importing basic libraries import os +import numpy as np import supervision as sv import wget import torch @@ -173,18 +174,18 @@ def single_image_detection(self, img, img_path=None, det_conf_thres=0.2, id_stri return self.results_generation([lab, box, scrs], img_path, id_strip) - def batch_image_detection(self, data_path, batch_size=16, det_conf_thres=0.2, id_strip=None): + def batch_image_detection(self, data_source, batch_size=16, det_conf_thres=0.2, id_strip=None): """ Perform detection on a batch of images. - + Args: - data_path (str): - Path containing all images for inference. + data_source (str or List[np.ndarray]): + Either path containing images for inference or list of numpy arrays (RGB format, shape: H×W×3). batch_size (int, optional): Batch size for inference. Defaults to 16. - det_conf_thres (float, optional): + det_conf_thres (float, optional): Confidence threshold for predictions. Defaults to 0.2. - id_strip (str, optional): + id_strip (str, optional): Characters to strip from img_id. Defaults to None. extension (str, optional): Image extension to search for. Defaults to "JPG" @@ -192,14 +193,20 @@ def batch_image_detection(self, data_path, batch_size=16, det_conf_thres=0.2, id Returns: list: List of detection results for all images. """ - dataset = pw_data.DetectionImageFolder( - data_path, - transform=self.transform, - ) - + # Handle numpy array input + if isinstance(data_source, (list, np.ndarray)): + image_source = ((Image.fromarray(np.asarray(img)).convert('RGB'), str(i)) + for i, img in enumerate(data_source)) + # Handle image directory input + else: + dataset = pw_data.DetectionImageFolder( + data_source, + transform=self.transform, + ) + image_source = ((Image.open(img_path).convert('RGB'), img_path) for img_path in dataset.images) + results = [] - for i in range(len(dataset)): - im_pil = Image.open(dataset.images[i]).convert('RGB') + for im_pil, img_id in image_source: w, h = im_pil.size orig_size = torch.tensor([w, h])[None].to(self.device) im_data = self.transform(im_pil)[None].to(self.device) @@ -210,8 +217,8 @@ def batch_image_detection(self, data_path, batch_size=16, det_conf_thres=0.2, id lab = labels[0][scr > det_conf_thres] box = boxes[0][scr > det_conf_thres] scrs = scores[0][scr > det_conf_thres] - - res = self.results_generation([lab, box, scrs], dataset.images[i], id_strip) + + res = self.results_generation([lab, box, scrs], img_id, id_strip) # Normalize the coordinates for timelapse compatibility size = orig_size[0].cpu().numpy() diff --git a/PytorchWildlife/models/detection/ultralytics_based/yolov5_base.py b/PytorchWildlife/models/detection/ultralytics_based/yolov5_base.py index 917b6d93c..5a23bd4ed 100644 --- a/PytorchWildlife/models/detection/ultralytics_based/yolov5_base.py +++ b/PytorchWildlife/models/detection/ultralytics_based/yolov5_base.py @@ -133,12 +133,12 @@ def single_image_detection(self, img, img_path=None, det_conf_thres=0.2, id_stri return res - def batch_image_detection(self, data_path, batch_size: int = 16, det_conf_thres: float = 0.2, id_strip: str = None) -> list[dict]: + def batch_image_detection(self, data_source, batch_size: int = 16, det_conf_thres: float = 0.2, id_strip: str = None) -> list[dict]: """ Perform detection on a batch of images. Args: - data_path (str): Path containing all images for inference. + data_source (str or List[np.ndarray]): Either path containing images for inference or list of numpy arrays (RGB format, shape: H×W×3). batch_size (int, optional): Batch size for inference. Defaults to 16. det_conf_thres (float, optional): Confidence threshold for predictions. Defaults to 0.2. id_strip (str, optional): Characters to strip from img_id. Defaults to None. @@ -147,8 +147,33 @@ def batch_image_detection(self, data_path, batch_size: int = 16, det_conf_thres: list[dict]: List of detection results for all images. """ + # Handle numpy array input + if isinstance(data_source, (list, np.ndarray)): + results = [] + num_batches = (len(data_source) + batch_size - 1) // batch_size # Calculate total batches + + with tqdm(total=num_batches) as pbar: + for start_idx in range(0, len(data_source), batch_size): + batch_arrays = data_source[start_idx:start_idx + batch_size] + imgs = torch.stack([self.transform(img) for img in batch_arrays]).to(self.device) + predictions = self.model(imgs)[0].detach().cpu() + predictions = non_max_suppression(predictions, conf_thres=det_conf_thres) + + for idx, pred in enumerate(predictions): + pred = pred.numpy() + # Get size directly from numpy array + size = batch_arrays[idx].shape[:2] + pred[:, :4] = scale_boxes([self.IMAGE_SIZE] * 2, pred[:, :4], size).round() + res = self.results_generation(pred, f"{start_idx + idx}", id_strip) + # Normalize the coordinates for timelapse compatibility + res["normalized_coords"] = [[x1 / size[1], y1 / size[0], x2 / size[1], y2 / size[0]] for x1, y1, x2, y2 in pred[:, :4]] + results.append(res) + pbar.update(1) + return results + + # Handle image directory input dataset = pw_data.DetectionImageFolder( - data_path, + data_source, transform=self.transform, )