import cv2
from ultralytics import YOLO
import time


def detect_bikes_modern(video_path, output_path=None):
    # 1. 加载模型
    # yolov8n.pt 是 "Nano" 版本，速度最快，适合 CPU 运行
    # 第一次运行时会自动从互联网下载该文件 (约 6MB)
    print("正在加载 YOLOv8 模型...")
    model = YOLO('yolov8n.pt')

    # 2. 打开视频
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"错误: 无法打开视频 {video_path}")
        return

    # 获取视频参数
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # 3. 设置保存 (可选)
    out = None
    if output_path:
        # 现代环境通常使用 mp4v 或 avc1
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        print(f"结果将保存到: {output_path}")

    print("开始检测... 按 'q' 退出")

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

        # 4. 核心检测逻辑
        # stream=True 让处理更流畅
        # classes=[1] 意思是我们只检测 COCO 数据集中的第 1 类 (0是人, 1是自行车, 2是汽车...)
        # conf=0.3 意思是置信度大于 0.3 才算检测到
        results = model.predict(frame, conf=0.2, classes=[1], verbose=False)

        # 5. 在图上画框
        # result.plot() 会自动把检测框画在 frame 上
        annotated_frame = results[0].plot()

        # 显示计数
        bike_count = len(results[0].boxes)
        cv2.putText(annotated_frame, f"Bikes: {bike_count}", (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        # 显示画面
        cv2.imshow("Modern Bike Detection", annotated_frame)

        # 保存画面
        if out:
            out.write(annotated_frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    if out:
        out.release()
    cv2.destroyAllWindows()
    print("处理完成。")


if __name__ == "__main__":
    # 替换成你的视频路径
    # 注意：这里我们不再需要那个 lowerFAR.xml 文件了，直接扔掉它
    video_source = "intersection.avi"
    output_file = "result_modern.mp4"

    detect_bikes_modern(video_source, output_file)