main.py

import argparse
from pathlib import Path

from ultralytics import YOLO


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run YOLOv8s logo detection on a single image."
    )
    parser.add_argument(
        "--weights",
        type=str,
        default="best.pt",
        help=(
            "Path to the YOLOv8s weights file (e.g. best.pt) or a Hugging Face repo id "
            "(e.g. lakshya-rawat/yolov8s-pdf-logo-detector)."
        ),
    )
    parser.add_argument(
        "--source",
        type=str,
        required=True,
        help="Path to the input image (e.g. a rendered first page of a PDF).",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="If set, display the image with detections overlaid.",
    )
    parser.add_argument(
        "--save",
        action="store_true",
        help="If set, save the annotated image next to the source file.",
    )
    return parser.parse_args()


def run_inference(weights: str, source: str, show: bool = False, save: bool = False) -> None:
    # Load model (from local file or Hugging Face repo id)
    model = YOLO(weights)

    # Run inference
    results = model(source)

    # Print simple, reusable summary of detections
    print(f"Detections for {source}:")
    for r in results:
        for box in r.boxes:
            xyxy = box.xyxy[0].tolist()
            conf = float(box.conf[0])
            cls_id = int(box.cls[0])
            print(
                f" - class_id={cls_id}, confidence={conf:.3f}, "
                f"bbox=[{xyxy[0]:.1f}, {xyxy[1]:.1f}, {xyxy[2]:.1f}, {xyxy[3]:.1f}]"
            )

        if show:
            r.show()

        if save:
            # Save annotated image next to the source
            src_path = Path(source)
            out_path = src_path.with_name(src_path.stem + "_detections" + src_path.suffix)
            r.save(filename=str(out_path))
            print(f"Annotated image saved to: {out_path}")


def main() -> None:
    args = parse_args()

    if not args.weights:
        raise SystemExit("Error: --weights must be provided (path or repo id).")

    if not args.source:
        raise SystemExit("Error: --source must be provided (path to image).")

    run_inference(weights=args.weights, source=args.source, show=args.show, save=args.save)


if __name__ == "__main__":
    main()