事前学習済みモデルを使って表情認識をやってみようと思います。
はじめに
人間の判断とそれほど相違ない結果になってそうでした。
使うもの
- Yolov5(顔の検出モデル)
- Multi-task EfficientNet-B2(表情認識モデル)
Multi-task EfficientNet-B2とは?
Multi-taskとありますが、これは単一のモデルに複数のタスクを学習させることを言います。
Multi-task EfficientNet-B2では、VGG Face2データセットで事前にトレーニングされたFace recognition CNNを用いて、性別・年齢・人種を分類するようにマルチタスク学習を行います。
〇トレーニングのデータセット
- 性別・年齢予測ではIMDB-WikiとAdienceデータセット
- 人種予測ではUTKFaceデータセット
〇モデル
学習した性別・年齢・人種は、同一人物であれば、画像間の変化は少ないのに対し、顔の表情は、同一人物であっても大きく異なります。
そのため、Face recognition CNNで抽出した特徴量をそのまま表情認識には使用できないと考え、表情認識のために、ネットワーク全体を表情認識データセットであるAffectNetを用いてFine Tuningしています。
実装のポイント解説
全体の実装は、こちらにあります。
①Yolov5で顔の検出
②検出された画像に対してMulti-task EfficientNet-B2を用いて表情認識
表情認識のラベルは以下
(※8クラス分類になってます。)
idx_to_class = {
0: '怒り', # Anger
1: '軽蔑', # Contempt
2: '嫌悪感', # Disgust
3: '恐れ', # Fear
4: '幸福', # Happiness
5: '普通', # Neutral
6: '悲しみ', # Sadness
7: '驚き' # Surprise
}
yolo_weight = "/content/face-emotion-recognition/yolov5_face_detection/weights/best.pt"
yolo_dir = "/content/face-emotion-recognition/yolov5_face_detection/yolov5"
d_model = torch.hub.load(yolo_dir, 'custom', path=yolo_weight, source='local')
# 画像の読み込み
img = cv2.imread(input_imgs[0])
# 顔検出_yolov5
pred_im = d_model(img)
# 画像のcrop
det_img = pred_im.crop(save=False)
# 1枚の画像に複数の顔が検出されるのでfor文を使用する
for idx in range(len(det_img)):
# Yoloの推論結果から顔画像取得
face_img = det_img[idx]["im"]
表情認識を行う
MODEL_PATH = '/content/face-emotion-recognition/models/affectnet_emotions/enet_b2_8.pt'
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = torch.load(MODEL_PATH)
model=model.to(DEVICE)
model.eval()
img_transforms = transforms.Compose(
[
transforms.Resize((260,260)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
]
)
# 前処理
img_tensor = img_transforms(Image.fromarray(face_img))
img_tensor.unsqueeze_(0)
# モデル推論
scores = model(img_tensor.to(DEVICE))
# 予測結果
scores=scores[0].data.cpu().numpy()
# 表情認識結果
text = idx_to_class[np.argmax(scores)]
print("感情: ", text)
たったこれだけで表情認識できるなんてすごいですね。
まとめ
これを使って合コン中に相手の表情を取得してやろうと思います!(^^)!
他の事前学習済みモデルを使用すると年齢や性別についても予測できるみたいです。
論文には、Multi-task EfficientNet-B2は、モバイルアプリやエッジデバイスで動作するように軽量かつ高精度で認識することを目的に構築されたそうなのでIoT端末でもリアルタイムに動作するかも検証してみようと思います。