🤓

顔画像の品質を評価できるMagFaceをローカルGPU環境で動かしてみた

2024/12/14に公開

はじめに

以前、顔認識（Face Recognition）について調べたときに気になっていた「MagFace」という顔認識モデルが「MoZuMa」というライブラリに実装されていたので、ローカルのGPU環境で動かしてみました。

https://arxiv.org/abs/2103.06627

MagFaceとは？

MagFaceは、ArcFaceなどと同じく顔認識（顔特徴量抽出）モデルの1つで、「顔画像の品質」（顔の識別し易さ）を特徴量の大きさとして扱うことができるのが特徴です。
詳しくは、以下の日本語記事を参照ください。

MoZuMaとは？

MoZuMaは、様々な機械学習モデルを手軽に使えるようにしたライブラリ、みたいです。今回、初めて触りました。

https://github.com/mozuma/mozuma

上記のページでは、以下の機能が挙げられています。

Text to image retrieval: テキストによる画像検索
Image similarity search: 画像による画像検索（類似画像検索）
Image classification: 画像分類
Face detection: 顔検出
Object detection: 物体検出
Video keyframes extraction: 動画キーフレーム抽出
Multilingual text search: 多言語テキスト検索

今回は顔検出機能の1つとして実装されているMagFaceを使ってみます。

ローカルGPU環境でMagFaceを動かす

環境構築

今回は以下の環境で試しました。

GPU: NVIDIA GeForce RTX 4080 16GB
OS: Ubuntu 22.04.5 LTS
NVIDIAドライバ: 550.127.08
Docker: 24.0.7
Docker Compose: 2.18.0
NVIDIA Container Toolkit: 1.13.5
Python: 3.12.8

簡単に再現できるようにcompose.yaml、Dockerfile、requirements.txtを作りました。
以下のコマンドでPythonのREPL環境を起動できます。なお、ビルド後のイメージサイズは約5.3GBです。

docker-compose run --build --rm magface

compose.yaml

services:
  magface:
    build: "."
    deploy:
      resources:
        reservations:
          devices:
            - driver: "nvidia"
              capabilities: ["gpu"]
              count: 1

Dockerfile

FROM python:3.12-slim
WORKDIR /app
COPY ./requirements.txt ./
RUN --mount=type=cache,mode=0755,sharing=locked,target=/root/.cache/pip \
  python -m pip install --requirement requirements.txt

requirements.txt

boto3==1.35.81
botocore==1.35.81
certifi==2024.8.30
charset-normalizer==3.4.0
click==8.1.7
dill==0.3.9
facenet-pytorch==2.6.0
filelock==3.16.1
fsspec==2024.10.0
ftfy==6.3.1
huggingface-hub==0.26.5
idna==3.10
Jinja2==3.1.4
jmespath==1.0.1
joblib==1.4.2
MarkupSafe==3.0.2
mozuma==0.9.0
mozuma-clip==1.0.1
mpmath==1.3.0
networkx==3.4.2
numpy==1.26.4
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.19.3
nvidia-nvjitlink-cu12==12.6.85
nvidia-nvtx-cu12==12.1.105
opencv-python-headless==4.10.0.84
packaging==24.2
pillow==10.2.0
python-dateutil==2.9.0.post0
pytorch-ignite==0.5.1
PyYAML==6.0.2
regex==2024.11.6
requests==2.32.3
s3transfer==0.10.4
scikit-learn==1.6.0
scipy==1.14.1
six==1.17.0
sympy==1.13.3
threadpoolctl==3.5.0
tokenizers==0.21.0
torch==2.2.2
torchvision==0.17.2
tqdm==4.67.1
typing_extensions==4.12.2
urllib3==2.2.3
wcwidth==0.2.13
yacs==0.1.8

なお、手動で環境を構築する際には、以下のPythonパッケージをインストールする必要があります。

pip install mozuma

CUDAランタイムが適切に依存関係に含まれているため、CUDA Toolkitなどは不要です。

顔検出＆顔特徴量抽出

顔検出、顔特徴量抽出する例は以下の通りです。
顔検出モデルとしてMTCNN、顔特徴量抽出モデルとしてMagFaceをそれぞれ使っています。

なお、画像は「ぱくたそ」からお借りしました。

import torch
from mozuma.callbacks.memory import (
    CollectBoundingBoxesInMemory,
    CollectFeaturesInMemory,
)
from mozuma.models.magface.pretrained import torch_magface
from mozuma.models.mtcnn.pretrained import torch_mtcnn
from mozuma.torch.datasets import (
    ImageBoundingBoxDataset,
    ImageDataset,
    LocalBinaryFilesDataset,
)
from mozuma.torch.options import TorchRunnerOptions
from mozuma.torch.runners import TorchInferenceRunner
import requests

# 画像を取得する
# REF: https://www.pakutaso.com/20240534144post-51319.html
url = "https://user0514.cdnw.net/shared/img/thumb/minnadewaiwaiHFKE1900_TP_V4.jpg?&download=1"
with open("pakutaso_91037.jpg", "wb") as file:
    file.write(requests.get(url).content)

# モデルを読み込む
device = torch.device("cuda")
mtcnn = torch_mtcnn(device=device)
magface = torch_magface(device=device)

# 顔検出する
detection_dataset = ImageDataset(LocalBinaryFilesDataset(["pakutaso_91037.jpg"]))
bbox_memory = CollectBoundingBoxesInMemory()
runner = TorchInferenceRunner(
    model=mtcnn,
    dataset=detection_dataset,
    callbacks=[bbox_memory],
    options=TorchRunnerOptions(
        data_loader_options={"batch_size": 1}, device=device
    ),
)
runner.run()

# 顔特徴量抽出する
extraction_dataset = ImageBoundingBoxDataset(
    image_dataset=ImageDataset(LocalBinaryFilesDataset(bbox_memory.indices)),
    bounding_boxes=bbox_memory.bounding_boxes,
)
feature_memory = CollectFeaturesInMemory()
runner = TorchInferenceRunner(
    model=magface,
    dataset=extraction_dataset,
    callbacks=[feature_memory],
    options=TorchRunnerOptions(
        data_loader_options={"batch_size": 3}, device=device
    ),
)
runner.run()

顔検出の結果はbbox_memory、顔特徴量抽出の結果はfeature_memoryにぞれぞれ格納されています。顔特徴量は512次元です。

>>> bbox_memory
CollectBoundingBoxesInMemory(indices=['pakutaso_91037.jpg'], bounding_boxes=[BatchBoundingBoxesPrediction(bounding_boxes=array([[469.66794 , 247.95665 , 557.7657  , 371.4424  ],
       [260.236   , 226.52773 , 325.19904 , 314.0591  ],
       [115.475   , 137.14598 , 177.0021  , 213.08894 ],
       [676.06995 ,  59.889378, 731.61694 , 137.27159 ],
       [384.83875 ,  15.363367, 435.8606  ,  83.58076 ]], dtype=float32), scores=array([0.9314614 , 0.9999491 , 0.9997973 , 0.99994516, 0.9999664 ],
      dtype=float32), features=array([[[502.10785 , 287.14297 ],
        [531.81244 , 279.60278 ],
        [519.4442  , 308.44467 ],
        [509.82214 , 334.11502 ],
        [534.43976 , 326.44418 ]],

       [[271.72867 , 258.8479  ],
        [300.95343 , 259.4327  ],
        [282.9691  , 280.6257  ],
        [276.51883 , 292.3363  ],
        [299.9389  , 292.5765  ]],

       [[147.92133 , 166.44894 ],
        [163.93462 , 158.91643 ],
        [172.4172  , 171.1022  ],
        [162.67055 , 193.68187 ],
        [176.08365 , 186.9336  ]],

       [[685.67725 ,  85.26222 ],
        [706.9709  ,  88.85768 ],
        [683.95795 , 100.97431 ],
        [680.4966  , 118.73818 ],
        [695.3959  , 121.3723  ]],

       [[406.3634  ,  40.114075],
        [428.96902 ,  40.867157],
        [421.3183  ,  55.04886 ],
        [405.23    ,  64.380135],
        [426.6399  ,  64.81298 ]]], dtype=float32))])
>>> feature_memory
CollectFeaturesInMemory(indices=[('pakutaso_91037.jpg', 0), ('pakutaso_91037.jpg', 1), ('pakutaso_91037.jpg', 2), ('pakutaso_91037.jpg', 3), ('pakutaso_91037.jpg', 4)], features=array([[ 1.0638033 ,  0.6281913 , -0.26310876, ..., -1.957676  ,
        -0.8814139 , -0.24361716],
       [ 0.11199856,  0.95571923, -0.16173913, ..., -0.98713946,
         2.2939835 ,  0.18821453],
       [ 1.2645271 ,  0.22915061,  0.63601696, ..., -1.7950549 ,
         0.01394585,  0.04445985],
       [ 0.48207766,  0.58846843, -0.7248775 , ..., -0.17341338,
        -0.9320125 ,  0.46618274],
       [-0.15281442, -0.25264448, -0.6024233 , ...,  0.07988756,
        -0.115779  ,  1.1148782 ]], dtype=float32))
>>> feature_memory.features.shape
(5, 512)

ちゃんと5人検出されているようです。結果を画像にオーバーレイしてみましょう。

import numpy as np
from PIL import Image, ImageDraw

image = Image.open("pakutaso_91037.jpg")
draw = ImageDraw.Draw(image)

# バウンディングボックスを描画する
for x1, y1, x2, y2 in bbox_memory.bounding_boxes[0].bounding_boxes:
    draw.rectangle([x1, y1, x2, y2], outline="red", width=2)

# 顔器官点（キーポイント）を描画する
for key_points in bbox_memory.bounding_boxes[0].features:
    for x, y in key_points:
        draw.ellipse((x - 2, y - 2, x + 2, y + 2), fill="blue")

# スコア、マグニチュードを描画する
scores = bbox_memory.bounding_boxes[0].scores
features = feature_memory.features
magnitudes = np.linalg.norm(features, axis=1)
for i in range(len(bbox_memory.bounding_boxes[0].bounding_boxes)):
    x1, y1, x2, y2 = bbox_memory.bounding_boxes[0].bounding_boxes[i]
    draw.text(
        (x1, y1 - 16),
        f"s:{scores[i]:.3f} / m:{magnitudes[i]:.3f}",
        fill="green", font_size=16,
    )

image.save("pakutaso_91037.png")

オーバーレイ結果は以下の通りです。

右下の女性は顔を手で覆っているせいか顔検出、顔器官点（キーポイント）検出の結果がイマイチですが、それ以外の男性は妥当な検出結果のようです。
また、正面を向いている人ほどマグニチュード（顔画像の品質、顔の識別し易さ）が大きく、顔認識する時には役立ちそうです。

おわりに

MoZuMaを使うことで簡単にMagFaceを使った顔特徴量抽出を行うことができました。
今回の実行例ではマグニチュードの違いはそれほど分かりやすくありませんでしたが、多くの画像のマグニチュードを比べてみたところそれなりに違いは出たので、使える指標になりそうです。

本記事が何らかの参考になれば幸いです。