本記事では、DifFaceと呼ばれる機械学習手法を用いて任意の画像を超解像する方法をご紹介します。

出典: zsyOAOA/DifFace

DifFace

概要

DifFaceは、Diffusion modelベースのblind face restoration技術です。

blind face restorationとは、低解像やボケなどの未知の劣化を持つ低解像度の顔画像から、高品質の顔画像を復元するタスクを指します。

従来技術では、トレーニングデータに複雑な劣化が含まれている場合、復元精度が低下する問題や、忠実度・知覚的損失・敵対的損失などの複数の制約が必要であり、これらの影響を安定させるために、煩雑なハイパーパラメータの調整が必要であるという問題がありました。

DifFaceでは、低品質画像(LQ)から高品質画像(HQ)への事後分布(posterior)を確立することにより、複雑な損失設計なしでロバストな復元を実現しています。
LQ画像から事前トレーニング済み拡散モデルの中間状態への遷移分布を設計し、モデルを再帰的に適用することで、中間状態からHQ画像へと徐々に遷移させます。この方法により、SOTAを達成しています。

出典: DifFace: Blind Face Restoration with Diffused Error Contraction

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、任意の画像を復元していきます。

デモ(Colaboratory)

それでは、実際に動かしながら任意の画像を復元していきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

初めにGithubからソースコードを取得します。

%cd /content

!git clone https://github.com/zsyOAOA/DifFace.git

%cd /content/DifFace
# Commits on Jan 19, 2023
!git checkout 35d91a1c029b8a910b5a66684e0b51862bbe8e34

次にライブラリをインストールします。

%cd /content/DifFace


!pip install einops==0.6.0 omegaconf==2.3.0 timm==0.6.12

最後にライブラリをインポートします。

%cd /content/DifFace

import cv2
import matplotlib.pyplot as plt
from pathlib import Path

以上で環境セットアップは完了です。

Utility関数の定義

ここでは、出力結果の画像表示用のUtility関数を定義します。

def display(img1, img2):
  fig = plt.figure(figsize=(25, 10))
  ax1 = fig.add_subplot(1, 2, 1) 
  plt.title('Input', fontsize=16)
  ax1.axis('off')
  ax2 = fig.add_subplot(1, 2, 2)
  plt.title('DifFace', fontsize=16)
  ax2.axis('off')
  ax1.imshow(img1)
  ax2.imshow(img2)

def imread(img_path):
  img = cv2.imread(img_path)
  if img.ndim > 3:
    img = img[:, :, :3]
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  return img

Face Image Restoration

それでは、Blind Face Restorationを実施しています。

まず、顔部分のみを切り出した画像を復元していきます。
以下では、./testdata/cropped_facesに保存されているテスト画像を復元しています。

%cd /content/DifFace

!python inference_difface.py \
  --aligned \
  --in_path ./testdata/cropped_faces \
  --out_path ./result_aligned \
  --gpu_id 0

復元結果を表示して確認します。

%cd /content/DifFace

in_dir = "./testdata/cropped_faces"
im_path_list = sorted([x for x in Path(in_dir).glob("*.png")])
im_path_list.extend([x for x in Path(in_dir).glob("*.jpg")])
out_dir = Path("./result_aligned/restored_faces")

for im_path_in in im_path_list:
  im1 = imread(str(im_path_in))
  im_path_out = out_dir / im_path_in.name
  im2 = imread(str(im_path_out))
  
  display(im1, im2)

出力結果は以下の通りです。

Whole Image Enhancement

続いて、画像全体の復元を行っていきます。

以下では、./testdata/whole_imgsに格納された画像を復元しています。

%cd /content/DifFace

!python inference_difface.py \
  --in_path ./testdata/whole_imgs \
  --out_path ./result_unaligned \
  --gpu_id 0

出力結果は以下の通りです。

%cd /content/DifFace

in_dir = "./testdata/whole_imgs"
im_path_list = sorted([x for x in Path(in_dir).glob("*.png")])
im_path_list.extend([x for x in Path(in_dir).glob("*.jpg")])
out_dir = Path("./result_unaligned/restored_image")

for im_path_in in im_path_list:
  im1 = imread(str(im_path_in))
  im_path_out = out_dir / im_path_in.name
  im2 = imread(str(im_path_out))
  
  display(im1, im2)

まとめ

本記事では、DifFaceを用いたBlind Face Restorationを行う方法をご紹介しました。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - DifFace: Blind Face Restoration with Diffused Error Contraction

2. GitHub - zsyOAOA/DifFace