运用yolov5人脸识别模型生成标准的一寸照视频

前段时间导师布置了一个人脸识别一寸照片的任务,给大家顺便分享一下如何使用yolov5(v6.1)训练好的face模型应用到简单的视频一寸照裁剪上。

我们可以把任务拆分成四步:
一:用widerface数据集在yolov5上训练出模型
二:将待测的视频逐帧转化为图片
三:用人脸识别模型对图片识别,并且进行一寸照人脸的裁剪,保存为新的图片
四:把生成的图片再转换为视频

一:用widerface数据集在yolov5上训练出模型

训练模型用的是widerface数据集,将数据集转换成YOLO格式之后,修改对应的配置文件,就可以开始训练了。这里我用我自己训练好的模型best.pt来进行测试。

二:将待测的视频逐帧转化为图片

首先我们要将一段视频拆分成图片,视频的帧率就是每秒播放多少张图片。
我找了一段有14秒的夸张人脸的视频,定为face.mp4 ,视频效果如下:

face

运用yolov5人脸识别模型生成标准的一寸照视频_第1张图片

然后用以下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帧/秒
运用yolov5人脸识别模型生成标准的一寸照视频_第2张图片
视频转化后的图片如下,均为jpg格式:
运用yolov5人脸识别模型生成标准的一寸照视频_第3张图片

三:用人脸识别模型对图片识别,并且进行一寸照人脸的裁剪,保存为新的图片

这个时候就需要用到我们训练好的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的下面配置文件处进行修改:
运用yolov5人脸识别模型生成标准的一寸照视频_第4张图片
只需要修改训练好的模型pt文件、需要识别的人脸图片保存的路径和yaml配置文件就可以了。

顺便贴上我的face.yaml文件的配置:
运用yolov5人脸识别模型生成标准的一寸照视频_第5张图片
由于只需要训练人脸一个目标,因此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


运用yolov5人脸识别模型生成标准的一寸照视频_第6张图片
视频中每一帧图片的大小都是295*413,并且符合头部占照片尺寸的2/3的标准。至此,我们便将一个动态的视频通过静态的YOLOv5检测方法,转化为人脸一寸照的视频格式了。

你可能感兴趣的:(python,计算机视觉,pytorch)