一、准备数据
先准备一个MP4的视频
1.测试一帧
使用opencv来提取每一个视频的帧
先使用以下代码查看一帧的内容,是否符合预期
import cv2
import matplotlib.pyplot as plt# 打开视频文件
video = cv2.VideoCapture("111.mp4")
# 读取一帧
ret, frame = video.read()plt.imshow(frame)# 颜色通道转换为RBG
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
plt.imshow(frame_rgb)
plt.axis('off')
plt.show()
运行结果:
2.提取每一帧
video = cv2.VideoCapture("111.mp4")
num = 0 # 计数器
save_step = 10 # 间隔帧
while True:ret, frame = video.read()if not ret:breaknum += 1if num % save_step == 0:cv2.imwrite("./images/" + str(num) + ".jpg", frame)
运行结果:
3.标注工具
在之前创建的yolov5中安装labelimg
pip install labelimg
在终端中使用命令打开窗口
labelimg
在labelimg窗口中,点击"Open Dir",选择刚才提取视频帧的文件
选择"YOLO"格式,点击图中的位置会切换
点击"View"选择"Auto Save mode"
点击窗口中的"Save",选择保存路径
为每一帧图片进行标注
点击"Create RectBox",为图中人物进行框选,并命名,之后点击"ok",之后的每一张类似这样的操作
标注完成后,点击保存图片信息的文件夹
二、数据调整
- images:存放图片
- train:训练集图片
- val:验证集图片
- label:存放标签
- train:训练集标签文件(要与训练集图片名称一致)
- val:验证机标签文件(要与训练集图片名称一致)
创建一个dataset将数据全部存入其中
注意:label中的clesses.txt要在此文件的根目录下
最后的目录结构
三、模型训练参数设置
将images中的train和label,复制到yolov5文件夹下的data/images中
在data根目录下创建一个yaml文件
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ./datasets # dataset root dir
train: images/train # 训练集的图片
val: images/val # 验证集的图片
test: # test images (optional)# Classes.txt中的内容填入其中
names:0: baiyi1: xiaoqiao
将yolov5根目录下的train.py中修改
四、启动模型训练
在终端中使用命令启动
python train.py
正在运行的输出:
运行成功后:
点击runs/train/exp 可以看到结果
使用detect.py来检测目标视频
python detect.py --weights runs/train/exp/weights/best.pt --source dataset/111.mp4 --view-img
在runs/train/exp中即可找到目标文件