OpenCVで教師データの作成

main.py

# coding=utf-8
import sys

import numpy as np
import cv2

def split_image(filename):
    cascade_file="./lbpcascade_animeface.xml"
    cascade = cv2.CascadeClassifier(cascade_file)
    if cascade.empty():
        raise Exception('cascade file not found')

    cap = cv2.VideoCapture(filename)

    if not cap.isOpened():
        raise RuntimeError("{}: not found".format(filename))

    i = 0
    frame_count = 0
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret == True and frame_count % 2 == 0:
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            gray = cv2.equalizeHist(gray)
            faces = cascade.detectMultiScale(
                gray,
                scaleFactor = 1.1,
                minNeighbors = 5,
                minSize = (80, 80)
            )
            for (x, y, w, h) in faces:
                dst = frame[y:y+h, x:x+w]
                output_path = 'data/' + str(i) + '.png'
                print(output_path)
                cv2.imwrite(output_path, dst)
                i = i + 1

            cv2.imshow('frame', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        elif ret == True:
            pass
        else:
            break

        frame_count += 1

    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    if len(sys.argv) < 2:
        sys.stderr.write('error')
        sys.exit(-1)
    split_image(sys.argv[1])