diff --git a/scripts/run_inference.py b/scripts/run_inference.py index fe873ee..574f7b3 100644 --- a/scripts/run_inference.py +++ b/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: - - //_predicted.jpg (bbox overlay) - - //.json (mothra-annotator format) +Runs a trained YOLO model on a directory of images and writes into --output-dir: + - /_predicted.jpg (bbox overlay) + - /.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, )