基于yoloV5对视频进行目标检测并保存

发布于 2021-08-30  1745 次阅读


代码运行效果如视频所示:

https://www.bilibili.com/video/BV1xf4y1N7mF

只需将待处理视频所在文件夹、待处理视频路径、保存视频路径,以及原视频帧速率和像素大小设置好即可。

ps: 第一次运行需加载一段时间的模型;

import os
import cv2
import torch
import numpy as np

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

path = r'C:\Users\pbh0612\Desktop'
old_path = path + '\\test.mp4'
new_path = path + '\\m.mp4'
tem_path = path + '\\image0.jpg'# 临时保存图片的地方(无需修改,且保存的图片会自动删除)
# 设置处理后的视频保存时的参数(和原视频一致即可)
fps = 25    # 待保存视频的帧速率
size = (1920, 1080) # 待保存视频的size

# 加载原视频
cap = cv2.VideoCapture(old_path)
# 提前设置处理后视频的参数
fourcc = cv2.VideoWriter_fourcc(*'XVID')
videoWriter = cv2.VideoWriter(new_path, fourcc, fps, size)#第一个是要保存的地址,最后一个是保存图片的尺寸

while(cap.isOpened()):
    ret, frame = cap.read()
    if ret==True:
        if hasattr(torch.cuda, 'empty_cache'):# 释放内存,防止内存不足
            torch.cuda.empty_cache()# 释放内存,防止内存不足
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)# 更改BGR为RBG
        results = model([frame], size=640)# yolo训练传进来的帧图片
        results.save(path)# 将训练好的图片临时保存
        videoWriter.write(cv2.imread(tem_path))# 将训练好的帧图片写入新视频中
        os.remove(tem_path)
    else:
        break
        
videoWriter.release()

最后的声音可以直接用moviepy模块添加,博主为了方便,直接使用的pr软件添加。