Skip to content

Feat #576: accept NumPy array input in batch_image_detection for MDv5 and RT-DETR - #660

Open
proofbyhuman wants to merge 1 commit into
microsoft:mainfrom
proofbyhuman:feature/numpy-batch-detection-input
Open

Feat #576: accept NumPy array input in batch_image_detection for MDv5 and RT-DETR#660
proofbyhuman wants to merge 1 commit into
microsoft:mainfrom
proofbyhuman:feature/numpy-batch-detection-input

Conversation

@proofbyhuman

Copy link
Copy Markdown

Cuerpo del PR

Summary

batch_image_detection already accepts a list of NumPy arrays in YOLOV8Base, but the other detector
backends still require images on disk. This PR brings the same capability to YOLOV5Base
(MegaDetectorV5) and RTDETRApacheBase (MegaDetectorV6 Apache), so in-memory pipelines — images stored
as binary columns in Parquet, PySpark/Dask workflows — no longer need to write intermediate JPEG files
just to run detection.

Closes #576.

Changes made

PytorchWildlife/models/detection/ultralytics_based/yolov5_base.py

  • Added a # Handle numpy array input branch that mirrors the existing one in yolov8_base.py:
    batches the arrays, stacks them with the same MegaDetector_v5_Transform already used by
    single_image_detection, runs NMS with det_conf_thres, and rescales boxes with scale_boxes
    using each array's own shape[:2].
  • Renamed the first parameter data_pathdata_source to match the YOLOV8Base signature.
  • The directory branch is untouched.

PytorchWildlife/models/detection/rtdetr_apache/rtdetr_apache_base.py

  • This method already iterated one image at a time, so instead of duplicating the loop I made the image
    source a lazy generator: NumPy arrays become PIL images via Image.fromarray, and the directory case
    keeps using DetectionImageFolder and opens each file inside the loop exactly as before (no eager
    loading of the whole folder).
  • Same data_pathdata_source rename.

In both backends, array inputs get the index as img_id ("0", "1", …) and one result entry per
input image, matching what YOLOV8Base already does.

Why this works

The transforms in both backends already accepted in-memory images — single_image_detection has taken
ndarray input in yolov5_base.py and rtdetr_apache_base.py for a long time. The only thing missing
was a path into the batch method, so this reuses the existing preprocessing rather than adding a second
code path for it.

Testing

Tested locally on CPU with the real pretrained weights, comparing both input modes over the same four
images. The images are deliberately of different, non-square sizes — 800×600, 640×480, 1024×768 and one
portrait 600×800 — so that any width/height mix-up in the box rescaling would show up as a mismatch:

  • MegaDetectorV5 (md_v5a.0.0.pt): ran batch_image_detection(folder), then
    batch_image_detection(list_of_arrays) with the same images loaded via
    np.array(Image.open(p).convert("RGB")). Both passes returned 4 results with 16 detections in total,
    and xyxy, confidence, class_id and normalized_coords matched exactly.
  • MegaDetectorV6 Apache (MDV6-apa-rtdetr-c.pth): same comparison, same outcome — identical
    detections from disk and from arrays.
  • Both runs also assert that the array branch returns one result per input image.

Out of scope / notes for maintainers

  1. yolo_mit_base.py is not included. It sets cfg.task.data.source to the path and reloads the
    model on every call, so array support there is a different change. Happy to send a follow-up PR.
  2. data_pathdata_source rename. No caller inside the repo passes it as a keyword (all are
    positional), and it aligns the three backends, but it is technically a breaking change for anyone
    calling batch_image_detection(data_path=...) externally. Tell me if you'd rather keep the old name.
  3. Empty-detection images in MDv5. The existing directory branch does if pred.size(0) == 0: continue,
    so images with no detections are dropped from the results, while yolov8_base.py keeps them. The new
    array branch keeps them (consistent with YOLOv8). I did not change the directory branch, but this
    inconsistency looks related to Mismatch in number of results between detector and classifier #610 and I can open a separate issue/PR if useful.
  4. Pre-existing bug in rtdetr_apache_base.py (unchanged by this PR, but confirmed while testing).
    normalized_coords is computed from size = torch.tensor([w, h]), so x / size[1] divides x by the
    height and y / size[0] divides y by the width. On the non-square test images the normalised values
    came out as high as 1.34 — exactly the 800/600 aspect ratio — instead of staying within [0, 1].
    This affects the Timelapse-compatible output for any non-square image. I kept the behaviour untouched
    so this PR stays focused; glad to fix it in a separate PR if you want.

AI assistance

Author: @proofbyhuman, with AI assistance from Claude Code (Claude Fable 5) for exploring the codebase
and drafting the change. I reviewed and tested everything myself, and the implementation deliberately
follows the pattern already present in yolov8_base.py rather than introducing a new one.

Thanks a lot for taking the time to review this — happy to rework any part of it.

…n 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 microsoft#576

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support NumPy Array Input for batch_image_detection

1 participant