feat: add batch inference mode to run_inference.py

scripts/run_inference.py

@@ -1,9 +1,15 @@
 """
-Batch inference script for Mothra.
+Inference script for Mothra.
 
-Runs a trained YOLOv8 model on a directory of images and writes:
-  - <output_dir>/<image_stem>/<image_stem>_predicted.jpg  (bbox overlay)
-  - <output_dir>/<image_stem>/<image_stem>.json           (mothra-annotator format)
+Runs a trained YOLO model on a directory of images and writes into --output-dir:
+  - <image_stem>/<image_stem>_predicted.jpg  (bbox overlay)
+  - <image_stem>/<image_stem>.json           (mothra-annotator format)
+  - all_predictions.json                     (aggregated)
+
+Two modes:
+  - Single (default): one image fed to the model at a time.
+  - Batch: images fed in chunks of --batch-size, letting YOLO process
+    multiple images per forward pass (faster on GPU).
 
 The JSON format matches https://github.com/DDMAL/mothra-annotator so outputs
 can be loaded directly into the annotator for correction.
@@ -41,12 +47,52 @@ def yolo_box_to_annotator_bbox(box_xywhn, img_w: int, img_h: int) -> list[int]:
     return [round(x), round(y), round(w * img_w), round(h * img_h)]
 
 
+def _build_session(image_path: Path, result, timestamp: str) -> dict:
+    """Build a mothra-annotator session dict from a single YOLO result."""
+    img_h, img_w = result.orig_shape
+    annotations = []
+    for box in result.boxes:
+        yolo_class_id = int(box.cls.item())
+        annotator_class_id = YOLO_TO_ANNOTATOR_CLASS.get(yolo_class_id, yolo_class_id + 1)
+        bbox = yolo_box_to_annotator_bbox(box.xywhn[0].tolist(), img_w, img_h)
+        annotations.append(
+            {
+                "id": str(uuid.uuid4()),
+                "classId": annotator_class_id,
+                "bbox": bbox,
+                "confidence": round(float(box.conf.item()), 4),
+                "timestamp": timestamp,
+            }
+        )
+    return {
+        "imageName": image_path.name,
+        "imageWidth": img_w,
+        "imageHeight": img_h,
+        "annotations": annotations,
+    }
+
+
+def _save_result(image_path: Path, result, session: dict, output_dir: Path) -> None:
+    """Write the overlay image and per-image JSON for one result."""
+    stem = image_path.stem
+    image_out_dir = output_dir / stem
+    image_out_dir.mkdir(parents=True, exist_ok=True)
+
+    overlay = result.plot()
+    cv2.imwrite(str(image_out_dir / f"{stem}_predicted.jpg"), overlay)
+    (image_out_dir / f"{stem}.json").write_text(json.dumps(session, indent=2))
+
+    n = len(session["annotations"])
+    print(f"  {image_path.name}: {n} detection{'s' if n != 1 else ''}")
+
+
 def run_inference(
     images_dir: Path,
     weights_path: Path,
     output_dir: Path,
     conf_threshold: float = 0.25,
     iou_threshold: float = 0.7,
+    batch_size: int = 1,
 ):
     model = YOLO(str(weights_path))
     image_paths = collect_images(images_dir)
@@ -55,61 +101,40 @@ def run_inference(
         print(f"No images found in {images_dir}")
         return
 
-    print(f"Found {len(image_paths)} images. Running inference...")
+    mode = "single" if batch_size == 1 else f"batch (size {batch_size})"
+    print(f"Found {len(image_paths)} images. Running inference [{mode}]...")
     output_dir.mkdir(parents=True, exist_ok=True)
 
     timestamp = datetime.now(timezone.utc).isoformat()
     all_predictions = []
 
-    for image_path in image_paths:
-        stem = image_path.stem
-        image_out_dir = output_dir / stem
-        image_out_dir.mkdir(parents=True, exist_ok=True)
-
-        results = model.predict(
-            source=str(image_path),
-            conf=conf_threshold,
-            iou=iou_threshold,
-            save=False,
-            verbose=False,
-        )
-
-        result = results[0]
-        img_h, img_w = result.orig_shape
-
-        # Save annotated image directly (avoids YOLO's project/runs/detect layout)
-        overlay = result.plot()  # BGR numpy array with boxes drawn
-        overlay_path = image_out_dir / f"{stem}_predicted.jpg"
-        cv2.imwrite(str(overlay_path), overlay)
-
-        annotations = []
-        for box in result.boxes:
-            yolo_class_id = int(box.cls.item())
-            annotator_class_id = YOLO_TO_ANNOTATOR_CLASS.get(yolo_class_id, yolo_class_id + 1)
-            bbox = yolo_box_to_annotator_bbox(box.xywhn[0].tolist(), img_w, img_h)
-            annotations.append(
-                {
-                    "id": str(uuid.uuid4()),
-                    "classId": annotator_class_id,
-                    "bbox": bbox,
-                    "confidence": round(float(box.conf.item()), 4),
-                    "timestamp": timestamp,
-                }
+    if batch_size == 1:
+        for image_path in image_paths:
+            results = model.predict(
+                source=str(image_path),
+                conf=conf_threshold,
+                iou=iou_threshold,
+                save=False,
+                verbose=False,
             )
-
-        session = {
-            "imageName": image_path.name,
-            "imageWidth": img_w,
-            "imageHeight": img_h,
-            "annotations": annotations,
-        }
-
-        json_path = image_out_dir / f"{stem}.json"
-        json_path.write_text(json.dumps(session, indent=2))
-
-        n = len(annotations)
-        print(f"  {image_path.name}: {n} detection{'s' if n != 1 else ''}")
-        all_predictions.append(session)
+            session = _build_session(image_path, results[0], timestamp)
+            _save_result(image_path, results[0], session, output_dir)
+            all_predictions.append(session)
+    else:
+        # Feed images in chunks so YOLO processes batch_size images per forward pass.
+        for chunk_start in range(0, len(image_paths), batch_size):
+            chunk = image_paths[chunk_start : chunk_start + batch_size]
+            results = model.predict(
+                source=[str(p) for p in chunk],
+                conf=conf_threshold,
+                iou=iou_threshold,
+                save=False,
+                verbose=False,
+            )
+            for image_path, result in zip(chunk, results):
+                session = _build_session(image_path, result, timestamp)
+                _save_result(image_path, result, session, output_dir)
+                all_predictions.append(session)
 
     all_json_path = output_dir / "all_predictions.json"
     all_json_path.write_text(json.dumps(all_predictions, indent=2))
@@ -118,7 +143,7 @@ def run_inference(
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Mothra batch inference with annotator-compatible JSON output")
+    parser = argparse.ArgumentParser(description="Mothra inference with annotator-compatible JSON output")
     parser.add_argument(
         "--images-dir",
         type=str,
@@ -139,6 +164,12 @@ def main():
     )
     parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
     parser.add_argument("--iou", type=float, default=0.7, help="IoU threshold for NMS")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Number of images per YOLO forward pass. 1 = single mode (default); >1 = batch mode.",
+    )
     args = parser.parse_args()
 
     run_inference(
@@ -147,6 +178,7 @@ def main():
         output_dir=Path(args.output_dir),
         conf_threshold=args.conf,
         iou_threshold=args.iou,
+        batch_size=args.batch_size,
     )