example_inference.py

"""
Example inference script for YOLO11n Text model.

Usage:
    python example_inference.py --image path/to/image.jpg
    python example_inference.py --image path/to/folder/
"""

import argparse
from pathlib import Path

from ultralytics import YOLO


def main():
    parser = argparse.ArgumentParser(description="Text Detection Inference")
    parser.add_argument("--image", required=True, help="Image file or folder path")
    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
    parser.add_argument("--iou", type=float, default=0.7, help="NMS IoU threshold")
    parser.add_argument("--imgsz", type=int, default=640, help="Input image size")
    parser.add_argument("--save", action="store_true", help="Save annotated images")
    parser.add_argument("--show", action="store_true", help="Display results")
    args = parser.parse_args()

    # Load model
    model_path = Path(__file__).parent / "best.pt"
    model = YOLO(str(model_path))

    # Run inference
    results = model.predict(
        source=args.image,
        conf=args.conf,
        iou=args.iou,
        imgsz=args.imgsz,
        save=args.save,
        show=args.show,
    )

    # Print results
    for result in results:
        print(f"\nImage: {result.path}")
        print(f"Shape: {result.orig_shape}")
        
        if result.boxes is not None and len(result.boxes) > 0:
            print(f"Detected {len(result.boxes)} text regions:")
            for i, box in enumerate(result.boxes):
                xyxy = box.xyxy[0].tolist()
                conf = box.conf[0].item()
                print(f"  [{i+1}] bbox: [{xyxy[0]:.1f}, {xyxy[1]:.1f}, {xyxy[2]:.1f}, {xyxy[3]:.1f}], conf: {conf:.3f}")
        else:
            print("No text regions detected")


if __name__ == "__main__":
    main()