Zenn
😀

Google Colabで高精度の顔検出モデル(RetinaFace)を試す

2020/12/27に公開
Python
#
深層学習
tech

概要

今回は、Google Colab+pytorchで顔検出のモデルを試してみました。

顔検出モデル

まず、顔検出のモデルを調べるところから始めます。
Paper with codeでface detectionのstate-of-the-artモデルを調べます。
https://paperswithcode.com/task/face-detection
face detectionのベンチマークとランキングがいろいろ出てきますが、ページの下のほうに、Githubのスター数が多いモデルが載っています。今回は、最もスター数の多い"RetinaFace: Single-stage Dense Face Localisation in the Wild" [1]を使うことにします。
本モデルの実装はたくさんありましたが、pytorch実装で、かつpipで簡単にインストールできるこのライブラリを使うことにしました。
https://github.com/ternaus/retinaface

Google colabで動かす

早速、Google colabでこのモデルを使います。
Google ColabでGPUを使ったpytorchによる画像認識については、以前の記事に書きました。

Githubのインストラクションに従って、まず、pipでライブラリをインストールします。

!pip install -U retinaface_pytorch

次に、必要なモジュールのインポートとモデルの準備を行います。学習済みモデルをダウンロードします。

import cv2
from matplotlib import pyplot as plt
from retinaface.pre_trained_models import get_model
from retinaface.utils import vis_annotations

model = get_model("resnet50_2020-07-20", max_size=2048)
model.eval()

モデルのダウンロードが終わったら、Colaboratoryに画像をアップロードします。ColaboratoryではローカルからファイルをアップロードするモジュールをNotebookに記述することができます。数枚の画像に対してテストするだけならこれで逐一アップロードすればいいので、今回はこれを使います。

from google.colab import files
files.upload()

「ファイル選択」ボタンを押して、ファイルをアップロードします。今回テストとして、pixabayから2つのフリーの画像を使うことにします。

画像がアップロードできたら、OpenCVで読み込みます。

path = 'workplace-1245776_1920.jpg'
image = cv2.imread(path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

pathはアップロードした画像の名前を使ってください。
今回使用するモデルで扱えるように、最後の行でBGR画像(OpenCVの標準形式)をRGB画像に変換しています。
次に、モデルによる推論を行います。

results = model.predict_jsons(image)
print(results)

推論時間は計算環境と画像の解像度によりますが、私の環境(Tesla T4, 1920x1440画像)では15秒程度かかりました。
Result:

[{'bbox': [284, 255, 385, 384], 'score': 0.999835729598999, 'landmarks': [[314, 307], [362, 308], [340, 328], [319, 351], [356, 353]]}, {'bbox': [1148, 255, 1252, 402], 'score': 0.9997597336769104, 'landmarks': [[1192, 313], [1235, 310], [1224, 343], [1201, 368], [1232, 366]]}, {'bbox': [225, 506, 342, 660], 'score': 0.99970942735672, 'landmarks': [[277, 570], [326, 565], [315, 596], [289, 627], [324, 623]]}, {'bbox': [1645, 152, 1759, 308], 'score': 0.9993517994880676, 'landmarks': [[1656, 222], [1689, 209], [1660, 247], [1677, 276], [1702, 265]]}, {'bbox': [943, 483, 1075, 661], 'score': 0.9988055229187012, 'landmarks': [[984, 558], [1044, 559], [1017, 590], [991, 623], [1035, 624]]}, {'bbox': [1387, 480, 1523, 735], 'score': 0.9084256887435913, 'landmarks': [[1416, 581], [1415, 589], [1390, 626], [1418, 675], [1415, 679]]}]

結果を見ると、検出された顔ごとにBounding boxの座標('bbox')、推論確率のスコア('score')、特徴点の座標('landmarks')が出力されていることがわかります。

検出した顔のbounding boxを描画してみます。

def draw_res(image,results):
    for r in results:
        bbox = r['bbox']
        if not bbox:continue
        cv2.rectangle(image,(bbox[0],bbox[1]),(bbox[2],bbox[3]),(0,0,255),thickness=10)
    return image

image_out = draw_res(image,results)
cv2.imwrite('out.jpg',image_out)

結果は検出結果はこんな感じになりました。


オクルージョンにもかなり強く、ごく一部でも写っていれば検出できていますね。

脚注

  1. Jiankang Deng, Jia Guo, Yuxiang Zhou, Jinke Yu, Irene Kotsia and Stefanos Zafeiriou. RetinaFace: Single-stage Dense Face Localisation in the Wild. arXiv:1905.00641, 2019. ↩︎

Discussion