作者:CSDN @ 养乐多
本文将介绍输入mp4视频数据,经过YOLO目标检测算法,输出带有目标框的新视频的代码脚本。
文章目录
一、代码
代码只需要安装好cv2、ultralytics,然后修改模型路径、输入mp4视频文件路径、保存文件的路径。点击运行即可进行对视频数据的推理,结果保存一个新的带有目标框的新视频。
python
import cv2
from ultralytics import YOLO
import datetime
# 加载YOLOv8n模型
model = YOLO('../预训练模型/yolov8n.pt') # 替换为你的模型路径
# 打开输入视频
input_video_path = '../测试mp4视频/yourInputVideo.mp4' # 替换为你的输入视频路径
cap = cv2.VideoCapture(input_video_path)
# 获取视频信息
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# 获取当前时间并生成时间戳后缀
now = datetime.datetime.now()
time_str = now.strftime("%Y%m%d_%H%M%S")
# 定义输出视频路径并加上时间戳后缀
output_video_path = f'./runsVideo/output_video_{time_str}.mp4' # 替换为你的输出视频路径
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
frame_number = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# 使用YOLOv8进行检测
results = model(frame)
# 解析结果并绘制检测框
for result in results:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0]) # 获取检测框坐标
confidence = float(box.conf[0]) # 获取置信度
cls = int(box.cls[0]) # 获取类别
label = model.names[cls] # 获取类别名称
# 绘制检测框和标签
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f"{label} {confidence:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# 获取当前时间并绘制到帧上
now = datetime.datetime.now()
time_str = now.strftime("%Y-%m-%d %H:%M:%S")
cv2.putText(frame, time_str, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
# 写入处理后的帧
out.write(frame)
frame_number += 1
print(f'Processing frame {frame_number}/{frame_count}')
# 释放资源
cap.release()
out.release()