Openface(二):顔の位置合わせを実現

最近オープンソースOpenfaceを学び、LFWデータセットの顔揃えを実現しました.これらの顔の位置合わせの画像は、後続の深さ学習ニューラルネットワークの訓練データセットである.コードは、一般的なpythonパケットに加えて、2つのローカルパケットを参照します.1つはdatasetで、画像generatorオブジェクトと補助出力を生成します.もう1つはdlib_alignは、前の記事でdlibライブラリを使って顔を揃えるバッグです.
このコードは、元のコードのparseをより簡潔なfireに変更し、Configクラスを使用して整列関数に必要な入力を管理します.実際に呼び出す場合はterminalに入力するだけです

python align.py alignMain --size=100

を行ないます.
コードおよびコメントは次のとおりです.

# align script
#   dlib  

import cv2
import numpy as np
import os
import sys
import random
import shutil
import fire

#          
fileDir = os.path.dirname(os.path.realpath(__file__))
#  flidDir       models   ，  dlib     
modelDir = os.path.join(fileDir, '..', 'models') 
dlibModelDir = os.path.join(modelDir, 'dlib') #    dilb          
myfaceDir = os.path.join(modelDir, 'myface') #             

sys.path.append(os.path.dirname(fileDir)) #               

import align.dataset as dataset
import align.dlib_align as dlib_align

class Config(object):
    inputDir = os.path.join(fileDir, '..', 'Images', 'Raw')
    outputDir = os.path.join(fileDir,'..', 'Images', 'Aligned')
    landmarks = 'outerEyesAndNose'
    dlibFacePredictor = os.path.join(dlibModelDir, 'shape_predictor_68_face_landmarks.dat')
    verbose = True
    size = 96
    skipMulti = False
    fallbackLfw = None

opt = Config()

def alignMain(**kwargs):
    for k,v in kwargs.items():
        setattr(opt,k,v)

    dataset.mkdirP(opt.outputDir) #           ，  ；  ，    。
    imgs = list(dataset.ImageIterator(opt.inputDir)) #       generator
    random.shuffle(imgs) #    imgs    
    landmarkMap = {
        'outerEyesAndNose': dlib_align.AlignDlib.OUTER_EYES_AND_NOSE,
        'innerEyesAndBottomLip': dlib_align.AlignDlib.INNER_EYES_AND_BOTTOM_LIP} #           ，    ，           ，            
    #                    ，  
    if opt.landmarks not in landmarkMap:
        raise Exception("Landmarks unrecognized: {} 
 Allowed landmarks: 
 outerEyesAndNose 
 innerEyesAndBottomLip 
".format(opt.landmarks))

    #          
    landmarkIndices = landmarkMap[opt.landmarks]

    alignFace = dlib_align.AlignDlib(opt.dlibFacePredictor)

    nFallbacks = 0
    for imgObject in imgs:
        print("=== {} ===".format(imgObject.path)) #imgObject.path Image    
        outDir = os.path.join(opt.outputDir, imgObject.cls) #imgObject.cls Image    ，                       
        dataset.mkdirP(outDir) #   outDir   ，  ；  ，    
        outputPrefix = os.path.join(outDir, imgObject.name) #          ，                
        imgName = outputPrefix + ".png" #       ，png  

        if os.path.isfile(imgName):
            if opt.verbose:
                print(" + Already found {}, skipping.".format(imgName))
        else:
            rgb = imgObject.getRGB()
            if rgb is None:
                if opt.verbose:
                    print(" + Unable to load {}".format(imgName))
                outRgb = None
            else:
                outRgb = alignFace.align(opt.size, rgb,
                                          landmark_indices=landmarkIndices,
                                          SkipMulti=opt.skipMulti) #      ，           
                if outRgb is None and opt.verbose:
                    print(" + Unable to align {}".format(imgName))
            if opt.fallbackLfw and outRgb is None:
                nFallbacks += 1
                deepFunneled = "{}/{}.jpg".format(os.path.join(opt.fallbackLfw, imgObject.cls), imgObject.name)
                shutil.copy(deepFunneled, "{}/{}.jpg".format(os.path.join(opt.outputDir, imgObject.cls), imgObject.name)) #  outRgb   ，  fallbackLfw  ，   fallbackLfw             
            if outRgb is not None:
                if opt.verbose:
                    print(" + Writing aligned file to disk.")
                outBgr = cv2.cvtColor(outRgb, cv2.COLOR_RGB2BGR)
                cv2.imwrite(imgName, outBgr) #              

    if opt.fallbackLfw:
        print('nFallbacks:', nFallbacks)


if __name__ == '__main__':
    fire.Fire()