File size: 732 Bytes
a3a407d | 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 | from ultralytics import YOLO
import cv2
import os
import argparse
def quick_test(weights="models/best.pt", img_path=None):
if img_path is None:
img_path = "dataset_subset/images/train/synth_train_0000.jpg"
if not os.path.exists(weights):
print(f"Error: Weights not found at {weights}")
return
if not os.path.exists(img_path):
print(f"Error: Image not found at {img_path}")
return
model = YOLO(weights)
results = model(img_path, conf=0.1)
for r in results:
print(f"Detections: {len(r.boxes)}")
for box in r.boxes:
print(f" - {model.names[int(box.cls[0])]}: {float(box.conf[0]):.2f}")
if __name__ == "__main__":
quick_test()
|