前段时间导师布置了一个人脸识别一寸照片的任务,给大家顺便分享一下如何使用yolov5(v6.1)训练好的face模型应用到简单的视频一寸照裁剪上。
我们可以把任务拆分成四步:
一:用widerface数据集在yolov5上训练出模型
二:将待测的视频逐帧转化为图片
三:用人脸识别模型对图片识别,并且进行一寸照人脸的裁剪,保存为新的图片
四:把生成的图片再转换为视频
训练模型用的是widerface数据集,将数据集转换成YOLO格式之后,修改对应的配置文件,就可以开始训练了。这里我用我自己训练好的模型best.pt来进行测试。
首先我们要将一段视频拆分成图片,视频的帧率就是每秒播放多少张图片。
我找了一段有14秒的夸张人脸的视频,定为face.mp4 ,视频效果如下:
face
然后用以下video2jpg.py脚本,将videopath下的MP4文件转化成out_path下的jpg图片
from cv2 import VideoCapture
from cv2 import imwrite
# 定义保存图片函数
# image:要保存的图片名字
# addr;图片地址与相片名字的前部分
# num: 相片,名字的后缀。int 类型
def save_image(image, addr, num):
address = addr + str(num) + '.jpg'
imwrite(address, image)
if __name__ == '__main__':
video_path = "/mnt/4T/ym2/projects/yolo-max/yolov5-master/data/videos/face.mp4" # 视频路径(修改)
out_path = "/mnt/4T/ym2/projects/yolo-max/yolov5-master/data/video2jpg/" # 保存图片路径+名字(修改)
is_all_frame = True # 是否取所有的帧
sta_frame = 1 # 开始帧
end_frame = 80 # 结束帧
######
time_interval = 1 # 时间间隔
# 读取视频文件
videoCapture = VideoCapture(video_path)
# 读帧
success, frame = videoCapture.read()
print(success)
i = 0
j = 0
if is_all_frame:
time_interval = 1
while success:
i = i + 1
if (i % time_interval == 0):
if is_all_frame == False:
if i >= sta_frame and i <= end_frame:
j = j + 1
print('save frame:', i)
save_image(frame, out_path, j)
elif i > end_frame:
break
else:
j = j + 1
print('save frame:', i)
save_image(frame, out_path, j)
success, frame = videoCapture.read()
运行完成之后,会显示视频转化成的图片数,这里的frame就是图片数,共有437张图片,视频时长大概是14秒左右,视频帧率约为30帧/秒
视频转化后的图片如下,均为jpg格式:
这个时候就需要用到我们训练好的YOLOv5人脸识别模型了,利用训练好的YOLOv5人脸识别模型,在detect.py中对代码进行修改,读取出识别框的四个a、b、c、d坐标值。下列代码对应的是detect.py的第168行,直接在下面加即可
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
annotator.box_label(xyxy, label, color=colors(c, True))
#xyxy中保存的就是识别框的四个坐标值,我们可以直接把它取出来用
修改后的代码如下:
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
annotator.box_label(xyxy, label, color=colors(c, True))
#xyxy中保存的就是识别框的四个坐标值,我们可以直接把它取出来用
# wrote by ym
a = xyxy[0].item()
b = xyxy[1].item()
c = xyxy[2].item()
d = xyxy[3].item()
rate=1.3
w=c-a
cw = int((a + c) / 2)
h=d-b
ch = int((d + b) / 2)
h = int(413 * w / 295)
box = (cw - rate * w, ch - rate * h, cw + rate * w, ch + rate * h)
#print(box) #box保存的是识别框的四个坐标
# save_path保存的的是处理的每张图片的名字和路径,例如 runs/detect/exp52/133.jpg
#print(save_dir) # 值为runs/detect/exp19
#p保存的是每张待处理图片的路径!
img=Image.open(str(p))
print(p)
ss = str(p).rsplit("\\",1)[1]
im=img.crop(box)
im2 = im.resize((295, 413),Image.ANTIALIAS)
if save_crop:
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
关于对人脸的一寸照裁剪,参考如下:https://aistudio.baidu.com/aistudio/projectdetail/438751
只要从中截取一寸照裁剪的代码即可
同时在detect.py的下面配置文件处进行修改:
只需要修改训练好的模型pt文件、需要识别的人脸图片保存的路径和yaml配置文件就可以了。
顺便贴上我的face.yaml文件的配置:
由于只需要训练人脸一个目标,因此nc设为1即可
最后,再通过jpg2video.py脚本,把jpg图片按顺序拼接成一段MP4视频,到此用yolov5实现视频中人脸的一寸照识别
"""
该脚本的功能:
——————————将imgpath下的jpg文件转化成videopath下的mp4文件
fps代表帧率,也就是一帧播放多少图片
im.size表示图片的长宽
foucc代表视频的编码格式
视频默认按照图片的名字顺序进行生成
"""
import cv2
import os
from PIL import Image
def PicToVideo(imgPath, videoPath):
images = os.listdir(imgPath) # 获取所有图片名称
images.sort(key=lambda x: int(x[:-4])) # 对图片进行排序
fps = 30 # 帧率
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # 设置视频的编码格式
im = Image.open(imgPath + images[0]) # 获取图片信息
# videoWriter = cv2.VideoWriter(videoPath, fourcc, fps, (295,413), isColor=True) # 创建一个videoWriter对象,用于将图片保存为视频的一部分
videoWriter = cv2.VideoWriter(videoPath, fourcc, fps, im.size, isColor=True) # 创建一个videoWriter对象,用于将图片保存为视频的一部分
for im_name in range(len(images)):
frame = cv2.imread(imgPath + images[im_name])
videoWriter.write(frame)
videoWriter.release() # 释放
imgPath = "/mnt/4T/ym2/projects/yolo-max/yolov5-master/data/video22jpg/" # 图片的路径(修改)
videoPath = "/mnt/4T/ym2/projects/yolo-max/yolov5-master/data/jpg2video/face_cut.mp4" # 保存视频的路径及视频(修改)
PicToVideo(imgPath, videoPath)
最后看一下效果:
face_cut3