けの〜のブログ

ガッキーでディープラーニングして教育界に革命を起こす

動画から画像を切り出す

extractframes.py

import cv2
import os

def extract_frames(input_file, output_directory):
    '''
    extract frames from a video and save them in jpg
    input_file: file path of a video ここに動画データを格納する 形式はmp4
    output_directory: file path where frames will be saved  画像の保存先をここに入れる
    '''
    # TODO: rewrite this!
    try:
        os.makedirs(output_directory)
    except:
        pass

    vidcap = cv2.VideoCapture(input_file)

    success,image = vidcap.read()
    count = 0
    success = True

    while success:
        print('reading a frame: ', count)
        success,image = vidcap.read()

        output_file_path = os.path.join(output_directory, '{}.jpg'.format(count))
        cv2.imwrite(output_file_path, image)

        count += 1

if __name__ == '__main__':
    extract_frames('videos/0.mp4', 'frames/0')
extractface.py

import cv2
import glob
import os

OUTPUT_IMAGE_SIZE = (128, 128)
# CASCADE_PATHは自分のどこに"haarcascade_frontalface_alt.xml"があるかのpathを通す
CASCADE_PATH = '/usr/local/opt/opencv/share/OpenCV/haarcascades/haarcascade_frontalface_alt.xml'
cascade = cv2.CascadeClassifier(CASCADE_PATH)

def extract_faces(input_directory, output_directory):
    '''
    extract faces from frames and save them in jpg (128 * 128)
    input_directory: file path to where frames are stored (will be created by extractframes.extract_frames())
    output_directory: file path where faces will be saved
    '''

    # TODO: rewrite this!
    try:
        os.makedirs(output_directory)
    except:
        pass

    frame_files = glob.glob('{}/*'.format(input_directory))
    cnt = 0

    for frame_file in frame_files:
        try:
            print('processing {}'.format(frame_file))

            image = cv2.imread(frame_file)
            image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            image_h, image_w = image.shape[:2]
            face_rect_min_size = int(min(image_h, image_w) / 4)
            facerect = cascade.detectMultiScale(image_gray, scaleFactor=1.1, minNeighbors=1, minSize=(face_rect_min_size, face_rect_min_size))

            if not len(facerect) == 1:
                continue

            rect = facerect[0]

            x, y = rect[0], rect[1]
            w, h = rect[2], rect[3]

            cv2.imwrite(os.path.join(output_directory, '{}.jpg'.format(cnt)), cv2.resize(image[y:y+h, x:x+w], OUTPUT_IMAGE_SIZE))
            cnt += 1
        except Exception as e:
            # TODO: rewrite this!
            print(e)
main.py

import extractface
import extractframes
from pathlib import Path
import glob

def get_faces_from_video(input_file):
    # TODO: replace this!
    name = input_file.replace('videos/', '').replace('.mp4', '')

    # get frames from video
    extractframes.extract_frames(input_file = 'videos/{}.mp4'.format(name), output_directory = 'frames/{}'.format(name))

    # get faces from frames
    extractface.extract_faces(input_directory = 'frames/{}'.format(name), output_directory = 'faces/{}'.format(name))


if __name__ == '__main__':
    for file_name in glob.glob('videos/*'):
        get_faces_from_video(file_name)