77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
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) |