📝

dlibのshapre predictorでAFLW2000の学習をさせる

2024/06/30に公開

shape predictorについて

dlibの顔ランドマーク推定などを擁する推論器の種類です。5点予測、68点予測のモデルが用意されています。68点については「HELENデータセット」で学習させられているようです。
ざっと論文を見ると、モデルは回帰木の組み合わせで、初期座標から徐々に被写体の顔に合わせた位置にランドマークの座標を更新していくという計算方法の様です。
この「初期座標」はどうやら訓練データに依存して決まるようで、データセットを変えたら推論結果もそこそこ変わるかも?と思いました。

モチベーション

そこで、AFLW20003Dデータセットで2次元のランドマーク推論器を学習させなおしてみることにしました。こちらのデータセットは斜め、横向きの顔も多いので、より対応できるケースが増えるかも、という期待です。

訓練方法

下記ページを参考にしました。
Training a custom dlib shape predictor

公式でtrainingコードは用意されており、データセットに対応したxmlファイルを用意すればよさそうでしたので、サンプルをみつつ用意しました。

xmlの構成は下のような感じ

<dataset>
  <images>
    <image file='ファイルパス' width='450' height='450'> #画像の情報
      <box top=......> #bounding boxの座標
        <part name='00' x='<x座標>' y='<y座標>' />
        ... #以下68個のランドマークの列挙
      </box>
      ...# 以下複数人移っている場合は繰り返し
    </image>
    ...# 以下画像ファイルごとに記述するの
  </images>
</dataset>

解釈は比較的わかりやすいですね。

今回用いるデータセットの場合、「imageXXX.mat」の形で正解ラベルの情報が用意されているので、プログラムから上記xmlファイルを生成するならscipyなどで読み出す必要があります。

画像の情報をxmlファイルに書き出す部分は以下のようにしました。

import scipy
import glob
# 前略. globとかでファイルpathを収集するなど
train_files = glob.glob('somewhere/*.jpg')


for file_path in train_files:
    img_name  = file_path.split('/')[-1]
    train_xmlfile.write(f"  <image file='{img_name}' width='450' height='450'>\n")
    data_name = img_name.split('.')[0]
    mat_file_path=file_path.split('.jp')[0] #ディレクトリ構造に依る
    mat = scipy.io.loadmat(f'{mat_file_path}.mat')
    pt3d = mat['pt3d_68']
    bbox_top = int(min(pt3d[1]))
    bbox_left = int(min(pt3d[0]))
    bbox_width = int(max(pt3d[0])) - bbox_left
    bbox_height = int(max(pt3d[1])) - bbox_top
    train_xmlfile.write(f"    <box top='{bbox_top}' left='{bbox_left}' width='{bbox_width}' height='{bbox_height}'>\n")
    for i in range(68):
        idx = str(i).zfill(2)
        train_xmlfile.write(f"      <part name='{idx}' x='{int(pt3d[0][i])}' y='{int(pt3d[1][i])}'/>\n")
    train_xmlfile.write("    </box>\n  </image>\n")

# 後略. xmlのfootterを書くなど

あとはdlibのリポジトリにある訓練コードを実行します。xmlファイルのパス指定部分の編集が必要です。
さらに言うと、そのまま実行すると学習完了後に指定したディレクトリにあるすべてのjpg画像について検出結果を表示するようになっているので、不要なら消してもいいかもしれません。

結果

私の環境では1h以上かかりました。2000枚のうち1400枚を学習に使用しています。

こんな感じでした。青い線が推論結果です。


ただ、やはり横顔などは全くうまくいきませんでした。そんな簡単な話ではないですね。

Discussion