PyTorchで顔のキーポイント座標回帰を最初から組んでKaggle提出までやってみた
はじめに
Kaggleの Getting Started コンペである
「Facial Keypoints Detection」を使って、
- PyTorchでDatasetを作る
- CNNで座標回帰する
- validationで評価する
- best modelを保存する
- submission.csvを作成して提出する
ところまで、一通り実装してみた。
昨今、画像のディープラーニングモデルも便利になってきて、YOLO、Segformerなどいろいろありますが、あえて、pytorchだけで簡単に実装できるシンプルなCNNからスタートしました。
このコンペは
「画像から意味のある座標点を推定する」
という処理を理解することで
これは以下のような実務にも応用できると感じました。
- 構造物画像の制御点検出
- ひび・損傷位置推定
- 医療画像ランドマーク検出
- 画像位置合わせ
- 3D復元やオルソ補正
使用コンペ
環境
- Kaggle Notebook
- Python
- PyTorch
- matplotlib
- pandas
- numpy
まずはデータ確認
CSVには、
- 顔画像
- 目・鼻・口などの座標
が含まれている。
画像は Image 列にスペース区切り文字列として格納されていた。
sample = train_df.iloc[0]
image = np.array(sample["Image"].split(), dtype=np.float32)
image = image.reshape(96, 96)
キーポイントを画像に重ねる
keypoints = sample.drop("Image").values.astype(np.float32)
xs = keypoints[0::2]
ys = keypoints[1::2]
plt.imshow(image, cmap="gray")
plt.scatter(xs, ys)
plt.show()
実際に表示すると、目・鼻・口の座標が対応していることを確認できた。
この確認はかなり重要で、
- x,y逆転
- 左右逆
- ラベルズレ
などがあると、学習自体が壊れる。
Dataset作成
class FacialKeypointDataset(Dataset):
def __init__(self, df):
self.df = df.dropna().reset_index(drop=True)
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
row = self.df.iloc[idx]
image = np.array(row["Image"].split(), dtype=np.float32)
image = image.reshape(96, 96)
image = image / 255.0
image = torch.tensor(
image,
dtype=torch.float32
).unsqueeze(0)
keypoints = row.drop("Image").values.astype(np.float32)
# 座標正規化
keypoints = keypoints / 96.0
keypoints = torch.tensor(
keypoints,
dtype=torch.float32
)
return image, keypoints
Dataset作成で意識したこと
今回は、まずパイプラインを最後まで通すことを優先し、欠損ラベルを含む行は dropna() で除外しました。
ただし、この方法では学習に使えるデータ数が減るため、本来は欠損していないキーポイントだけを使ってlossを計算するmask処理を入れた方がよい可能性があります。
また、画像は 255.0 で割って 0〜1 に正規化し、キーポイント座標も 96.0 で割って 0〜1 に正規化しました。
入力画像と出力座標のスケールをそろえることで、学習が安定しやすくなると考えたためです。
CNNモデル
まずはシンプルなCNNを作成。
今回は精度を最初から追い込みすぎるよりも、画像から座標を回帰する一連の流れを理解することを優先しました。
そのため、ResNetなどの事前学習済みモデルではなく、畳み込み層を数段重ねたシンプルなCNNから始めています。
まず小さい構成で動かすことで、Dataset、loss、validation、submission作成までの流れを確認しやすくなると考えました。
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.regressor = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 12 * 12, 512),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(512, 30)
)
def forward(self, x):
x = self.features(x)
x = self.regressor(x)
return x
train / validation 分割
train_part, val_part = train_test_split(
clean_df,
test_size=0.2,
random_state=42,
shuffle=True
)
今回は、学習中のモデルが未知データにどの程度対応できるかを確認するため、trainとvalidationに分割しました。
Kaggleのスコアだけを見ると、手元で改善したのか、たまたまPublic側に合ったのかが分かりにくいため、まずはvalidation lossを確認しながら進める形にしています。
ただし、今回は単純なランダム分割なので、より厳密に評価するならKFoldなども試す余地があります。
loss関数
今回は座標回帰問題として扱うため、損失関数には MSELoss を使用しました。
分類ではなく、各キーポイントの x, y 座標という連続値を予測するため、予測座標と正解座標の差を小さくする目的に合っていると考えました。
一方で、外れ値に強い SmoothL1Loss なども候補になるため、今後比較してみたい部分です。
validation loss が最良のモデルを保存
best_val_loss = float("inf")
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save(
model.state_dict(),
best_model_path
)
最後のepochではなく、
validationが最も良かった重み
を使う形にした。
学習結果
best_val_loss ≒ 0.00048
train / validation の動きも比較的安定していた。
特に座標を 0〜1 に正規化してから、validation loss が安定した。
lossの値だけでは、実際に顔のどの位置をどの程度予測できているか分かりにくいため、validation画像に予測結果を重ねて確認しました。
座標回帰では、数値上のlossだけでなく、目・鼻・口などの構造が自然に保たれているかを見ることも重要だと感じました。
prediction可視化
validation画像で予測結果を確認。
- 青 = 正解
- × = 推論結果
かなり顔の構造を維持したまま推定できていた。
submission.csv 作成
test.csvを推論し、
IdLookupTable.csv に合わせて submission を生成。
submission.to_csv(
"/kaggle/working/submission.csv",
index=False
)
スコア
Public Score : 3.59348
Private Score: 3.56800
シンプルCNNでもある程度戦える結果だった
今回やらなかったこと
今回は、まずPyTorchで一通りの流れを組み切ることを目的にしたため、以下は行っていません。
- 欠損ラベルを考慮したloss計算
- data augmentation
- ResNetなどの事前学習モデル
- heatmap方式
- KFold validation
- 学習率スケジューラの細かい調整
そのため、この記事ではスコア改善よりも「座標回帰パイプラインを最後まで通すこと」を重視しています。
次にやりたいこと
- 左右反転augmentation
- 明るさ・コントラスト調整
- 欠損ラベル対応
- ResNet backbone
- heatmap方式
- EarlyStopping
- Scheduler
などデータの前処理をしてからCNNを強化させる方針で行こうと思います
まとめ
PyTorchだけで、
- データ読み込み
- 可視化
- Dataset作成
- CNN学習
- validation
- best model保存
- 推論
- submission作成
まで、一通り組み切ることができた。
Getting Started コンペなので、期限などを気にせず、シンプルなCNNだとどうかなど色々ためせるので良いと思う。
また
「画像から座標を推定する」
という、色々なものに応用が利きそうな内容なのでよかった。
Notebook: