OpenCV SVMで顔認識を試してみた

はじめに

OpenCV SVMを使って、入力顔写真からトランプかプーチンかを判別する処理を実装して試してみましたので、備忘として残します。

開発環境

Ubuntu 18.04.4 LTS
Python 3.6.9
opencv 4.5.1
dlib 19.21.1

入力画像

こちらの方法で、トランプとプーチンの顔写真を収集しました。
学習用にトランプとプーチンの画像それぞれ１００枚を使いました。

まず結果から

テストデータとして別に分類しておいた、トランプ、プーチンそれぞれの画像１０枚の画像でテストしました。
以下の１０枚の画像の内、９枚を正しくトランプと認識しました。

プーチンと誤認識をしてしまった１枚は、上記赤枠内の画像です。
こちらは、やや横向きの画像となっていますので、機械学習では難しいのかもしれません。

プーチンについても以下の１０枚の画像の内、９枚を正しくプーチンと認識しました。

トランプと誤認識した１枚は、上記赤枠内の画像です。
どうも、少し若い頃の写真なような気がします。

学習処理

入力画像からSIFTで特徴量を抽出し、BoW VocabularyをKmeansで学習します。Vocabulary（クラスタ）数は40にしました。
Vocabularyを学習したら、BoW Vocabularyで表現した特徴量を抽出し、それを使ってSVMで学習します。

以下が該当ソースコードになります。

def train(file_manager, voc_file, svm_file):
    # create a sift and a flann
    sift = cv2.SIFT_create()
    flann = MyMatcher.createFlann()

    # create a BoW KMeans Trainer
    bow_trainer = MyBoW(sift, flann, BOW_NUM_CLUSTERS)

    # add samples to the trainer
    for i in range(BOW_NUM_TRAINING_SAMPLES_PER_CLASS):
        pos_path, neg_path = file_manager.getFile(i)
        bow_trainer.addSample(pos_path)
        bow_trainer.addSample(neg_path)

    # create vocabulary
    bow_trainer.createVocAndSave(voc_file)

    # prepare training data with BoW decriptors
    training_data = util.TrainingData()
    for i in range(SVM_NUM_TRAINING_SAMPLES_PER_CLASS):
        pos_path, neg_path = file_manager.getFile(i)
        pos_img = cv2.imread(pos_path, cv2.IMREAD_GRAYSCALE)
        neg_img = cv2.imread(neg_path, cv2.IMREAD_GRAYSCALE)
        pos_descriptors = bow_trainer.extractDescriptors(pos_img)
        neg_descriptors = bow_trainer.extractDescriptors(neg_img)
        training_data.set(pos_descriptors, 1)
        training_data.set(neg_descriptors, -1)

    # create a SVM and train it
    svm = MySVM()
    svm.config(100)
    data, labels = training_data.get()
    svm.train(data, labels, svm_file)

予測処理

予測はSVMのpredictを呼ぶだけです。学習と予測は別々に呼び出せるようにしたため、予測前に学習データの読み込みを行っています。

def predict(file_manager, voc_file, svm_file):
    # create a sift and a flann
    sift = cv2.SIFT_create()
    flann = MyMatcher.createFlann()

    # create a Bow KMeans Trainer and the load vocaabulary
    bow_trainer = MyBoW(sift, flann, BOW_NUM_CLUSTERS)
    bow_trainer.loadVoc(voc_file)

    # create a SVM and load the data
    svm = MySVM()
    svm.load(svm_file)

    # predict
    for i in range(file_manager.getNumFiles()):
        path = file_manager.getFile(i)
        img = cv2.imread(path)
        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        descriptors = bow_trainer.extractDescriptors(gray_img)
        result = svm.predict(descriptors)
        logger.debug('file={0}, result={1}'.format(path, result))

その他のソースコード

上記コードでは、OpenCVのオブジェクトをラップしたクラスを作って処理していますので、参考までにそのクラスのコードも下記に載せておきます。

BOW_NUM_TRAINING_SAMPLES_PER_CLASS = 30
SVM_NUM_TRAINING_SAMPLES_PER_CLASS = 100
BOW_NUM_CLUSTERS = 40

class MyBoW:
    def __init__(self, dextractor, dmatcher, cluster_count):
        self._dextractor = dextractor
        self._trainer = cv2.BOWKMeansTrainer(cluster_count)
        self._extractor = cv2.BOWImgDescriptorExtractor(dextractor, dmatcher)

    def addSample(self, path):
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        keypoints, descriptors = self._dextractor.detectAndCompute(img, None)
        if descriptors is None:
            logger.debug('No descriptor genearted')
        else:
            self._trainer.add(descriptors)

    def createVocAndSave(self, data_file=None):
        voc = self._trainer.cluster()
        self._extractor.setVocabulary(voc)
        if data_file is not None:
            with open(data_file, 'wb') as f:
                pickle.dump(voc, f)

    def loadVoc(self, data_file):
        if data_file is not None:
            with open(data_file, 'rb') as f:
                voc = pickle.load(f)
            self._extractor.setVocabulary(voc)

    def extractDescriptors(self, img):
        keypoints = self._dextractor.detect(img)
        return self._extractor.compute(img, keypoints)

class MySVM:
    def __init__(self):
        self._svm = cv2.ml.SVM_create()

    def train(self, training_data, training_labels, data_file):
        aa = np.array(training_data)
        self._svm.train(np.array(training_data), cv2.ml.ROW_SAMPLE,
                        np.array(training_labels))
        if data_file is not None:
            self._svm.save(data_file)

    def config(self, count):
        self._svm.setType(cv2.ml.SVM_C_SVC)
        self._svm.setKernel(cv2.ml.SVM_LINEAR)
        self._svm.setTermCriteria((cv2.TERM_CRITERIA_MAX_ITER, count, 1e-6))

    def load(self, data_file):
        new_svm = self._svm.load(data_file)
        self._svm = new_svm

    def predict(self, descriptros):
        return self._svm.predict(descriptros)

class MyMatcher:
    def createFlann():
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = {}
        return cv2.FlannBasedMatcher(index_params, search_params)

所感

confidence scoreも取得するように修正しようと思います。