# mask_test.py
import cv2
from ultralytics import YOLO

# 1. Model Load Karein
# Training ke baad aapki best file yahan hoti hai: runs/detect/train/weights/best.pt
# Agar aapki training 'train6' thi, toh path hoga: 'runs/detect/train6/weights/best.pt'
model_path = r'best.pt'

try:
    model = YOLO(model_path)
    print("Model successfully load ho gaya!")
except:
    print("Error: 'best.pt' nahi mili. Path check karein.")
    exit()

# 2. Camera On Karein
cap = cv2.pycam = cv2.VideoCapture(0)

print("Camera start ho raha hai... Band karne ke liye 'q' dabayein.")

while cap.isOpened():
    success, frame = cap.read()
    if not success:
        break

    # 3. Mask Detection (Model Prediction)
    results = model(frame, conf=0.5) # 0.5 matlab 50% sure hoga toh dikhayega

    # 4. Results dikhayein
    annotated_frame = results[0].plot()

    cv2.imshow("Face Mask Detection Test", annotated_frame)

    # 'q' dabane se window band ho jayegi
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()