File size: 1,867 Bytes
4b09a94 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
"""
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()
|