helmet_detection / test_image.py
saad1BM's picture
Upload 8 files
21a1685 verified
# test_image.py
from ultralytics import YOLO
import cv2
import os
# Training ke baad wala model path
model_path = 'runs/detect/train2/weights/best.pt'
if os.path.exists(model_path):
model = YOLO(model_path)
# Dataset ke test folder mein se kisi image ka path
# Aap manually koi bhi image ka path yahan paste kar sakte hain
img_path = input("Enter image path (e.g., test/images/1.jpg): ")
if os.path.exists(img_path):
results = model(img_path)
for r in results:
im_array = r.plot() # Detection draw karega
cv2.imshow('Final Test', im_array)
cv2.waitKey(0) # Press any key to close
cv2.destroyAllWindows()
else:
print("Image file nahi mili! Path check karein.")
else:
print("Model abhi tak train nahi hua. Training khatam hone ka intezar karein.")