import cv2
from ultralytics import YOLO

# 加载YOLOv8模型
model = YOLO(r'D:\1BNS_projects\yolo_safehat_num_v8\weights\best.pt')

# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 调整图像大小以适应模型输入要求
    resized_frame = cv2.resize(frame, (640, 480))

    # 在图像上运行目标检测
    # results = yolo.detect_image(resized_frame)
    results = model.predict(resized_frame, save=True)

    # 在图像上绘制检测结果
    # for result in results:
    #     label, confidence, bbox = result
    #     x, y, w, h = bbox
    #     cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
    #     cv2.putText(frame, f'{label}: {confidence:.2f}', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    # 显示结果
    cv2.imshow('YOLOv8 Video Detection', frame)

    # 按下q键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()