inference/detect_image.py

from ultralytics import YOLO
import cv2
import sys

# Check for command-line argument
if len(sys.argv) < 2:
    print("Usage: python detect_image.py <image_path>")
    sys.exit(1)

image_path = sys.argv[1]

# Load trained model
model = YOLO("best50.pt")

# Read image
frame = cv2.imread(image_path)
if frame is None:
    print(f"Error: Could not read image at {image_path}")
    sys.exit(1)

# Run detection with conf and iou threshold
results = model(frame, conf=0.5, iou=0.4, imgsz=640, augment=False)

# Plot and save results
annotated_img = results[0].plot()
cv2.imwrite("output.jpg", annotated_img)
cv2.imshow("Detection", annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()