Python + OpenCV で雑コラ動画を作成する④ 課題への対処

class FacePosition:
    '''
    顔の位置とIDをセットで保持するためのクラス
    IDと座標を顔の座標・サイズを持つただの構造体
    '''

    def __init__(self, id, coordinate):
        self.id = id
        self.coordinate = coordinate

class FaceFrame:
    '''
    各フレームで認識した顔を保持するためのクラス
    faceCountはアプリケーション全体でIDが被らないように使ったIDの数をカウントするための変数なので、
    使うときは必ずFaceFrame.faceCountでアクセスする
    '''
    
    faceCount = 0

    def __init__(self, frame, coordinates):

        '''
        フレームと認識した顔の座標・サイズを渡す。
        顔の数分だけFacePointクラスのインスタンスを作成
        coodinates: 顔認識結果の配列。cascade.detectMultiScaleの結果をそのまま渡す
        '''
        
        # 顔の数分配列を確保
        self.faces = [None]*len(coordinates)
        self.frame = frame

        # 渡された顔それぞれにidを割り振り、FacePositionのインスタンスを作成
        for i in range(0, len(coordinates)):
            self.faces[i] = FacePosition(FaceFrame.faceCount, coordinates[i])
            FaceFrame.faceCount += 1

    # フレーム内の顔を後から追加するための関数
    def append(self, faceId, coordinate):
        self.faces.append(FacePosition(faceId, coordinate))

class FrameManager:
    
    '''
    渡されたフレームと顔認識結果を元に顔の連続性と抜けた顔の補完を行うクラス
    連続している顔には同じIDを割り振る。
    '''
    #いくつのFaceFrameを元に顔の連続性を確認するかを指定
    LIST_SIZE = 5
    CENTER_INDEX = int(LIST_SIZE/2)
    # フレーム間の顔が同じ顔であるかを判断する際の、位置、サイズの差をどこまで認めるか。%で指定
    ALLOWED_GAP = 5
        
    def __init__(self, height, width):
        '''
        扱う動画の高さ、幅を指定する
        '''
        FrameManager.FRAME_HEIGHT = height
        FrameManager.FRAME_WIDTH = width

        self.__frames = [None]*self.LIST_SIZE
        
    
    def put(self, frame, coordinates):
        '''
        渡されたフレームと顔認識結果を元にフレームを追加する
        追加時にIDの割り振り、連続性確認、抜けた顔の補完をし、LIST_SIZE番目のFaceFrameのインスタンスを返す
        終了時の処理として、全てのフレームを処理し終えた後、LIST_SIZE個のフレームがFrameManager内に残るので、残ったフレームを出し終えるまで、Noneを追加し続けること。

        return: FaceFrameのインスタンス。但し、LIST_SIZE番目にFaceFrameインスタンスがない場合はNoneを返す。
        '''
        # 最後に残ったフレーム出力するときにNoneが渡されるので、その場合あfaceFrameもNoneとする。
        if frame is None:
            faceFrame = None
        else:
            faceFrame = FaceFrame(frame, coordinates)

        # リストを1つづつ前にずらし、最後尾に引数のフレームを追加する。内部処理にランダムアクセスが多いので配列で管理する方が望ましいと思う。
        returnFrame = self.__frames[0]
        for i in range(0,len(self.__frames)-1):
            self.__frames[i] = self.__frames[i+1]
        self.__frames[FrameManager.LIST_SIZE-1] = faceFrame

        # 前後のフレームから連続性を確認する
        # CENTER_INDEXを境にその前(i)後(j)それぞれの組み合わせで顔の連続性を確認する
        for i in range(0, FrameManager.CENTER_INDEX):
            for j in range(FrameManager.CENTER_INDEX+1, FrameManager.LIST_SIZE):
                # Noneの部分は飛ばす
                if self.__frames[i] is not None and self.__frames[FrameManager.CENTER_INDEX] is not None and self.__frames[j] is not None:

                    # 間にあるフレーム全てに連続性確認、補完を行う
                    for k in range(i+1, j):
                        self.connectFrame(self.__frames[i], self.__frames[k], self.__frames[j])

        return returnFrame
        
        
    def connectFrame(self, frameF, frameC, frameB):               
        # frameF.facesとframeC.facesで連続している顔があれば同じidを振る。　
        # TODO 同じidが複数の顔に振られうる可能性がある。そもそもこの場合だと今の設計ではうまくいかないので一旦保留。
        frontFaceNum = len(frameF.faces)
        centerFaceNum = len(frameC.faces)
        backFaceNum = len(frameB.faces)
        for i in range(0, frontFaceNum):
            # 前のフレームの中のi番目の顔がframeCの顔の中のどれかと合致したかを保持
            matched = False
            for j in range(0, centerFaceNum):
                # 同じ顔と判断したら同じIDにする
                if self.compare(frameF.faces[i], frameC.faces[j]) == True:
                    frameC.faces[j].id = frameF.faces[i].id
                    matched = True
                    break
                
            # frameCに無くてもframeFとframeBの両方にある場合は間のframCにもその顔があるとみなして補完する。
            if matched == False:
                for k in range(0, backFaceNum):
                    if self.compare(frameF.faces[i], frameB.faces[k]):
                        # frameFとframeBの中間の位置・サイズに顔を追加する
                        frameC.append(frameF.faces[i].id, ((frameF.faces[i].coordinate + frameB.faces[k].coordinate)/2).astype(np.int))
                        # 顔の数を1増やす。(あとの処理でもう1つ顔が見つかった場合のため)
                        centerFaceNum += 1

                        # 無限ループ防止
                        if(centerFaceNum>10):
                            break

       
    def compare(self, face1, face2):
        '''
        face1、face2が連続したものであるか比較する。
        return: 同じならTrue、違えばFalse
        '''
        result = True
        # 座標、顔のサイズの違いが誤差の範囲に収まるかを確認し、全て誤差(ALLOWED_GAP)内であれば同じ顔であると判断する
        # TODO フレーム間が離れているならその分許容誤差も大きくしたほうが良い
        for i in range(0,4):
            if i%2 == 0:
                gap = ((float(face1.coordinate[i])-float(face2.coordinate[i]))/FrameManager.FRAME_HEIGHT)*100
            else:
                gap = ((float(face1.coordinate[i])-float(face2.coordinate[i]))/FrameManager.FRAME_WIDTH)*100
            if (-1*FrameManager.ALLOWED_GAP < gap < FrameManager.ALLOWED_GAP) == False:
                result = False
                break
        return result

# -*- coding:utf-8 -*-

import cv2
import datetime
import numpy as np
from PIL import Image

import frame_manager

def overlay_movie2():

    # 入力する動画と出力パスを指定。
    target = "target/test_input.mp4"
    result = "result/test_output2.m4v"  #.m4vにしないとエラーが出る

    # 動画の読み込みと動画情報の取得
    movie = cv2.VideoCapture(target) 
    fps    = movie.get(cv2.CAP_PROP_FPS)
    height = movie.get(cv2.CAP_PROP_FRAME_HEIGHT)
    width  = movie.get(cv2.CAP_PROP_FRAME_WIDTH)

    # 形式はMP4Vを指定
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    
    # 出力先のファイルを開く
    out = cv2.VideoWriter(result, int(fourcc), fps, (int(width), int(height)))

    # カスケード分類器の特徴量を取得する
    cascade_path = "haarcascades/haarcascade_frontalface_alt.xml"
    cascade = cv2.CascadeClassifier(cascade_path)
    
    # FrameManagerの作成
    frameManager = frame_manager.FrameManager(height, width)

    #　認識した顔を囲む矩形の色を指定。ここでは白。
    color = (255, 255, 255) 
    
    # 最初の1フレームを読み込む
    if movie.isOpened() == True:
        ret,frame = movie.read()
    else:
        ret = False

    count = 0

    # フレームの読み込みに成功している間フレームを書き出し続ける
    while ret:
        
        # グレースケールに変換
        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # 顔認識の実行
        facerecog = cascade.detectMultiScale(frame_gray, scaleFactor=1.1, minNeighbors=1, minSize=(1, 1))
            
        # 認識した顔をFrameManagerに入れる
        managedFrame = frameManager.put(frame, facerecog)

        # 5回め以降はFrameManagerからフレームが返ってくるのでファイル出力
        if managedFrame is not None:

            # 認識した顔に番号を書き加える
            for i in range(0,len(managedFrame.faces)):

                # 扱いやすいように変数を用意
                tmpCoord = managedFrame.faces[i].coordinate
                tmpId = managedFrame.faces[i].id
                    
                print("認識した顔の数(ID) = "+str(tmpId))
                                        
                # 矩形で囲む
                cv2.rectangle(managedFrame.frame, tuple(tmpCoord[0:2]),tuple(tmpCoord[0:2]+tmpCoord[2:4]), color, thickness=2)
                    
                # 顔のIDを書き込み
                cv2.putText(managedFrame.frame,str(tmpId),(tmpCoord[0],tmpCoord[1]),cv2.FONT_HERSHEY_TRIPLEX, 2, (100,200,255), thickness=2)
    
            out.write(managedFrame.frame)
        if count%10 == 0:
            date = datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
            print(date + '現在フレーム数：'+str(count))
        
        count += 1
        ret,frame = movie.read()

        # 途中終了
        if count > 200 :
            break

    print("出力フレーム数："+str(count))

    
if __name__ == '__main__':
    overlay_movie2()