python通过手机拍摄的视频图片进行人脸头像采集

1、概要

在计算机视觉上,模型要想训练的好,必要的前提需要有大量的图片数据做支撑,我使用公开的人脸攻击视频数据集进行图片采集并进行训练,但是有一个问题需要处理,大量的视频数据集是手机拍摄的,但是我们使用opencv在采集视频每帧数据的时候有一个问题,由于不知道视频旋转方向导致后期的图片无法识别人脸头像进行训练,所以我们需要知道每个视频的旋转角度,然后将采集的每帧图片进行旋转矫正,用于后续的训练。

2、解决

我们通过ffmpeg模块可打印出手机拍摄的视频的信息,例如分辨率,大小,旋转角度等等信息,python中 通过调用info = ffmpeg.probe(file_name)方法可获取视频信息,视频的信息是一个json类型数据,我们可以在打印的信息中找到相关角度的信息,获取角度并实现对视频获取的图片进行翻转,同时我也对其进行了一个方法封装请往下看 。

{'streams': [{'index': 0, 'codec_name': 'h264', 'codec_long_name': 'H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10', 'profile': 'High', 'codec_type': 'video', 'codec_tag_string': 'avc1', 'codec_tag': '0x31637661', 'width': 480, 'height': 640, 'coded_width': 480, 'coded_height': 640, 'closed_captions': 0, 'has_b_frames': 2, 'sample_aspect_ratio': '1:1', 'display_aspect_ratio': '3:4', 'pix_fmt': 'yuv420p', 'level': 30, 'chroma_location': 'left', 'refs': 1, 'is_avc': 'true', 'nal_length_size': '4', 'r_frame_rate': '30/1', 'avg_frame_rate': '30/1', 'time_base': '1/15360', 'start_pts': 0, 'start_time': '0.000000', 'duration_ts': 182272, 'duration': '11.866667', 'bit_rate': '353103', 'bits_per_raw_sample': '8', 'nb_frames': '356', 'disposition': {'default': 1, 'dub': 0, 'original': 0, 'comment': 0, 'lyrics': 0, 'karaoke': 0, 'forced': 0, 'hearing_impaired': 0, 'visual_impaired': 0, 'clean_effects': 0, 'attached_pic': 0, 'timed_thumbnails': 0}, 'tags': { 'rotate': '270', 'language': 'eng', 'handler_name': 'VideoHandler', 'vendor_id': '[0][0][0][0]'}, 'side_data_list': [{'side_data_type': 'Display Matrix', 'displaymatrix': '\n00000000:            0      -65536           0\n00000001:        65536           0           0\n00000002:            0    31457280  1073741824\n', 'rotation': 90}]}, {'index': 1, 'codec_name': 'aac', 'codec_long_name': 'AAC (Advanced Audio Coding)', 'profile': 'LC', 'codec_type': 'audio', 'codec_tag_string': 'mp4a', 'codec_tag': '0x6134706d', 'sample_fmt': 'fltp', 'sample_rate': '48000', 'channels': 2, 'channel_layout': 'stereo', 'bits_per_sample': 0, 'r_frame_rate': '0/0', 'avg_frame_rate': '0/0', 'time_base': '1/48000', 'start_pts': 0, 'start_time': '0.000000', 'duration_ts': 586752, 'duration': '12.224000', 'bit_rate': '128097', 'nb_frames': '575', 'disposition': {'default': 1, 'dub': 0, 'original': 0, 'comment': 0, 'lyrics': 0, 'karaoke': 0, 'forced': 0, 'hearing_impaired': 0, 'visual_impaired': 0, 'clean_effects': 0, 'attached_pic': 0, 'timed_thumbnails': 0}, 'tags': {'language': 'eng', 'handler_name': 'SoundHandler', 'vendor_id': '[0][0][0][0]'}}], 'format': {'filename': './video/G_NT_HS_g_E_2_13.mp4', 'nb_streams': 2, 'nb_programs': 0, 'format_name': 'mov,mp4,m4a,3gp,3g2,mj2', 'format_long_name': 'QuickTime / MOV', 'start_time': '0.000000', 'duration': '12.258000', 'size': '734419', 'bit_rate': '479307', 'probe_score': 100, 'tags': {'major_brand': 'isom', 'minor_version': '512', 'compatible_brands': 'isomiso2avc1mp41', 'encoder': 'Lavf56.19.101'}}}

import ffmpeg

# 旋转角度获取
def get_video_rotate_angle(file_name):
    """ Get video rotate angle """
    try:
        info = ffmpeg.probe(file_name)
        return 360 - int(info['streams'][0]['tags']['rotate'])
    except KeyError as e:
        return 0
    except Exception as e:
        return -1

3、处理

解决了获取视频旋转角度的问题剩下的就能够剩下的问题便可迎刃而解,我们需要将视频按照每帧获取并在获取之前需要知道该视频是否是被旋转了的,并且通过上述封装方法获取旋转角度,使用imutils.rotate_bound(img,angle)进行图片翻转,使用dlib的68点定位获取每个图的人脸进行保存,为了提高执行效率,我们使用线程池对数据集进行处理,所有代码我放在下面。

# -*- coding: utf-8 -*-
""" 
@Time    : 2021/10/11 11:13
@Author  : liwei
@Description: 
"""
import cv2
import imutils
import ffmpeg
import os
import dlib
import time
import threadpool
# 旋转角度获取
def get_video_rotate_angle(file_name):
    """ Get video rotate angle """
    try:
        info = ffmpeg.probe(file_name)
        return 360 - int(info['streams'][0]['tags']['rotate'])
    except KeyError as e:
        return 0
    except Exception as e:
        return -1

# 保存图片
def save_image(image, addr):
    address = addr + '.jpg'
    cv2.imwrite(address, image)

# 获取视频每帧图片并旋转抠取人像
def get_video_image(videoName, savePath, angle, detector, predictor):
    # 读帧
    videoCapture = cv2.VideoCapture(videoName)
    date = time.strftime("%Y-%m-%d %H:%M:%S")
    # 计数
    i = 0
    print('{} save image start'.format(videoName))
    while True:
        success, img = videoCapture.read()
        try:
            if success:
                i = i + 1
                # save_image(img, savePath + "_" + str(i))
                # 旋转图片
                if angle > 0:
                    img = imutils.rotate_bound(img, angle)
                gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                dets = detector(gray, 1)
                for face in dets:
                    shape = predictor(img, face)
                    chang = []
                    kuan = []
                    # 遍历所有点,打印出其坐标,并圈出来
                    for pt in shape.parts():
                        chang.append(pt.x)
                        kuan.append(pt.y)
                    # 调节人像大小
                    x1, x2 = max(chang) + 60, min(chang) - 65
                    minSize = min(kuan)
                    y1, y2 = max(kuan) + 100, minSize - int(minSize/2)
                    cropped = img[y2 + 1:y1, x2 + 1:x1]
                    if cropped.shape[1] > 200:
                        cv2.imwrite(savePath + "_" + str(i) + '.jpg', cropped)
                        print(date + " " + savePath + "_" + str(i) + '.jpg')
                    #cropped = cv2.resize(cropped, (IMAGE_SIZE, IMAGE_SIZE), cv2.INTER_CUBIC)
            else:
                print('{} {} total {} images'.format(date, videoName, i))
                print('{} {} image save end'.format(date, videoName))
                break
        except Exception as e:
            continue
    videoCapture.release()
    cv2.destroyAllWindows()

def main(f):
    dirName = f.split("_")[0]
    videoName.add(dirName)

    # 不存在目录先创建
    videoDir = VIDEO_OUTPUT_IMAGE + dirName
    try:
        if not os.path.exists(videoDir):
            os.mkdir(videoDir)
    except FileExistsError:
        print("目录存在")

    # 获取视频旋转角度
    angle = get_video_rotate_angle(VIDEO_PATH + f)
    # 按帧截图
    get_video_image(VIDEO_PATH + f, videoDir + "/" + f, angle, detector, predictor)

if __name__ == '__main__':

    # 常量
    # TEST = "./test/"
    VIDEO_PATH = "./video/flv/"
    # VIDEO_PATH = "./test/"
    # VIDEO_OUTPUT_IMAGE = './output/face_image/'
    VIDEO_OUTPUT_IMAGE = './output/test/'
    FACE_MODEL = "./model/shape_predictor_68_face_landmarks.dat"
    IMAGE_SIZE = 64

    # 变量
    videoName = set()

    # 人脸分类器
    detector = dlib.get_frontal_face_detector()
    # 获取人脸检测器
    predictor = dlib.shape_predictor(
        FACE_MODEL
    )

    # 获取文件夹下视频
    videoFileList = os.listdir(VIDEO_PATH)
    # 开始计时
    start_time = time.time()
    pool = threadpool.ThreadPool(2)

    requests = threadpool.makeRequests(main, videoFileList)
    [pool.putRequest(req) for req in requests]
    pool.wait()
    # 结束时间
    end_time = time.time()
    print('%s' % (end_time - start_time))

Guess you like

Origin blog.csdn.net/m0_43432638/article/details/122472732