初めに

これをやろうと思ったきっかけがこちらのTweet

処理のざっとした流れとしては

物体検出で人物の位置を特定
顔認証で人物を特定し、特定の人物だった場合に画像を保存する

私立恵比寿中学の中山莉子さんを抜き出すことを出発点にしているが、実装にあたってはそれなりに汎用的に動くようにすることを目指した。

環境・ライブラリなど

環境

Windows 10 + WSL2 (Ubuntu 20.04.1 LTS)

CUDA対応済み。セットアップ方法は各々の環境に合わせていただければ...。
(GPU使う想定でしか動作確認してないです。)

主なライブラリ

Pytorch(1.7.1)
Open CV(4.4.0.46)
face_recognition (1.2.2)

物体検出にEfficientDetを使用したため
https://github.com/rwightman/efficientdet-pytorch
のmasterブランチをcloneして取り込んだ。

face_recognitionが依存しているdlibが公式にはMacとLinuxにしか対応していないようなので
注意が必要です。

物体検出

EfficientDetを使用した。実装にあたっては https://zenn.dev/nnabeyang/articles/30330c3ab7d449c66338 を参考にした。

Dataset

PytorchのDatasetクラスを継承する。継承にあたってはEfficientDetの実装にあるDetectionDatsetを参考にした。

最低限、__len()__ と __getitem__(index)を実装しておく必要があるので、以下のような実装にした。


import torch.utils.data as data
from pathlib import Path

class RikoDataset(data.Dataset):
    """
    Datasetの拡張
    """
    def __init__(self, image_paths, transform=None):
        super().__init__()

        self.image_paths = image_paths
        self._transform = transform

    def __len__(self):
        return len(self.image_paths)
    def __getitem__(self, index):
        img = Image.open(self.image_paths[index]).convert('RGB')
        target = dict(img_idx=index, img_size=(img.width, img.height))
        if self.transform is not None:
            img, target = self.transform(img, target)
    
        return img, target

    @property
    def transform(self):
        return self._transform

    @transform.setter
    def transform(self, t):
        self._transform = t

モデル構築

モデルを作成する。

あまり複雑にしない設定なので、以下のように実装した


from effdet import create_model

def riko_create_model():
    """
    Efficent Netのモデルを構築する
    """

    bench = create_model(
        "efficientdet_d1", # d0 ~ d7
        bench_task="predict",
        num_classes=20
    )

    return bench

Dataloaderの作成

EfficientDetが最初から持っている実装 create_loader を使用した。

    loader = create_loader(
            test_dataset,
            input_size=input_config['input_size'],
            batch_size=1,
            use_prefetcher=True,
            interpolation='bilinear',
            fill_color=input_config['fill_color'],
            mean=input_config['mean'],
            std=input_config['std'],
            num_workers=1,
            pin_mem=False)

切り出し候補の抽出

Dataset, model, Loaderが出来上がったので、あとは画像を読み込んで切り出し候補を抽出すればよい。
物体検出の役割はここまで。


    with torch.no_grad():
        for input_image, target in loader:
            target_image_path_index = target["img_idx"].item()
            image_path = image_paths[target_image_path_index]
            output = bench(input_image, target)[0]
            for i in range(output.size(0)):
                xmin, ymin, xmax, ymax, pred, label = output[i]
                #3,8, 11, 13,15,16がよさそう
                if label.item() == 8:
                    if image_path in result:
                        control_list = result[image_path]
                        control_list.append([xmin.item(),
                                            ymin.item(),
                                            xmax.item(),
                                            ymax.item()])
                        result[image_path] = control_list
                    else:
                        result[image_path] = [[xmin.item(),
                                                ymin.item(),
                                                xmax.item(),
                                                ymax.item()]]

関数化してまとめたらこうなった。


def get_bounding_box(image_paths):
    """
    受け取ったパスの画像から、人っぽいもののバウンディングボックスの
    座標を返す

    image_paths: 対象の画像のパスのリスト

    return
        切り出し候補の座標と画像のパス
    """

    bench = riko_create_model()
    bench = bench.cuda()
    bench.eval()

    test_dataset = RikoDataset(image_paths)
    input_config = resolve_input_config({}, bench.config)
    loader = create_loader(
            test_dataset,
            input_size=input_config['input_size'],
            batch_size=1,
            use_prefetcher=True,
            interpolation='bilinear',
            fill_color=input_config['fill_color'],
            mean=input_config['mean'],
            std=input_config['std'],
            num_workers=1,
            pin_mem=False)

    result = {}

    with torch.no_grad():
        for input_image, target in loader:
            target_image_path_index = target["img_idx"].item()
            image_path = image_paths[target_image_path_index]
            output = bench(input_image, target)[0]
            for i in range(output.size(0)):
                xmin, ymin, xmax, ymax, pred, label = output[i]
                #3,8, 11, 13,15,16がよさそう
                if label.item() == 8:
                    if image_path in result:
                        control_list = result[image_path]
                        control_list.append([xmin.item(),
                                            ymin.item(),
                                            xmax.item(),
                                            ymax.item()])
                        result[image_path] = control_list
                    else:
                        result[image_path] = [[xmin.item(),
                                                ymin.item(),
                                                xmax.item(),
                                                ymax.item()]]

    return result

顔認証

顔認証は、face_recognitionを使用した。

学習

どの写真に誰が写っているのか、学習させる必要がある。
写真と写っている人物を紐付ける安直な方法として

学習に使う写真と正解ラベルを書いたCSVを用意する。
pandasでロードしてDataFrameを作成する。

を考えたので、今回はそのようにした。以下のようなCSVファイルを作成する。

file_name,name
learn_mayama.png, mayama
learn_riko.png, riko

顔認証

基本的に、face_recognitionの使い方の内容をほぼなぞっている。

唯一の違いは、切り出したい画像をディスクから読むのではなく、Open CVの形式になっているものをPillow形式に変換し、ndarrayに変更していることである。

https://qiita.com/derodero24/items/f22c22b22451609908ee を参考にして、以下のような実装をした。

from PIL import Image
import cv2

def opencv2pillow(image):
    new_image = image.copy()
    new_image = Image.fromarray(new_image)
    return np.array(new_image)

全体的なつくりは以下のようになった。


import face_recognition
import cv2
import numpy as np
import glob
import pandas as pd


class FaceRecognition():
    """
    顔の識別を実施するクラス
    """

    def get_train_file_DataFrame(self, csv_path, target_path="../data/faces/"):
        """
        学習用の画像のパスを取得して、対応するファイルと人物のデータフレームを作成する

        csvのファイル形式(例)

        ```
        file_name,name
        1.jpg,riko
        2.jpg,mayama
        3.jpg,ayaka
        4.jpg,mirei
        5.jpg,hinata
        6.jpg,kaho
        ```
        """
        train_list = pd.read_csv(csv_path)
        train_list["file_name"] = train_list["file_name"].apply(lambda x: "{}{}".format(target_path, x))

        return train_list

    def detection_riko(self, df, target_image):
        """
        画像の情報が入っているデータフレームをもらってきて、その情報
        をもとに学習する
        """
        encoded_face_list = []
        for file_path in df["file_name"]:
            loaded_face_image = face_recognition.load_image_file(file_path)
            encoed_image = face_recognition.face_encodings(loaded_face_image)[0]

            encoded_face_list.append(encoed_image)

        face_locations = face_recognition.face_locations(target_image, model="cnn")
        face_encodings = face_recognition.face_encodings(target_image, face_locations)

        for face_encode in face_encodings:
            matches = face_recognition.compare_faces(encoded_face_list, face_encode)
            face_distances = face_recognition.face_distance(encoded_face_list, face_encode)

            best_match_index = np.argmin(face_distances)

            if matches[best_match_index]:
                return True
        
        return False

物体検出から顔認証、そして切り出し。

ここまで来れば、物体検出と顔認証をつなぎ合わせ、特定の人物を切り出して保存するだけになる。

def connect(image_paths, csv_path):
    """
    物体検出の結果を顔認証のエンジンに渡す

    image_paths: 対象の画像のパスのリスト
    csv_path: 訓練データリストの入っているcsvファイルのパス
    """
    target_image_imfos = get_bounding_box(image_paths)

    for image_path, candidate_list in target_image_imfos.items():
        loaded_image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        loaded_image = cv2.cvtColor(loaded_image, cv2.COLOR_BGR2RGB)
        candidate_images = []
        for image_points in candidate_list:
            # 小数点はだめなので、妥協の産物としてint型に変換しておく
            x = int(image_points[0])
            y = int(image_points[1])
            w = int(image_points[2])
            h = int(image_points[3])

            new_image = opencv2pillow(loaded_image[y:y + h, x:x + w])
            candidate_images.append(new_image)

        # 顔認証
        face_recognition = FaceRecognition()
        learn_df = face_recognition.get_train_file_DataFrame(csv_path, "data/faces/")

        for number, candidate_image in enumerate(candidate_images):
            compare_result = face_recognition.detection_riko(learn_df, candidate_image)

            if compare_result:
                save_name = "{}_result_{}.jpg".format(image_path, number)
                save_target_image = Image.fromarray(candidate_image)
                save_target_image.save(save_name, quality=95)
                print("{} saved.".format(save_name))
            else:
                print("Image not saved.")

問題点

最大の問題点は、同じ画像に対して物体検知を実施しているのに、毎回実行結果が異なる点。乱数的なものが裏で発生していると思っているのだが、そこまで深堀できてない。
labelが数値で返ってくるが、それが何を意味しているのか自分が分かっていない。
顔認証のところで再現できていない謎のエラーがある。(エラーメッセージを取りそこなった)

ソースなど

こちら。
https://github.com/Tomosato/riko_detection

疑問などあれば、issue書いてもらえるといいかと思います。
できるだけ対応します。