CodeFormerによる顔画像のハイクオリティな高画質化

この記事では、CodeFomerを用いて、低画質の顔画像をハイクオリティに高画質化する方法を紹介します。

CodeFormerとは
デモ
まとめ
参考文献

CodeFormerとは

CodeFormerは、低画質の画像から高画質の画像を復元するためのネットワークで、情報の多くが欠落している低画質画像から情報を復元するために、2段階の学習を行います。

1段階目は、高画質の画像を入力とし、あえて情報を落とした上で、その情報を復元するself-reconstructionを学習します。エンコーダーでベクトル化した各セルの情報を（四捨五入のようなイメージで）離散化して情報量を落とし、デコーダーで元の画像を復元します。ここで学習した離散化ルールとデコーダーは、2段階目に引き継がれます。

2段階目は、低画質画像から高画質画像の復元を学習します。低画質画像をエンコーダーでベクトル化した後、1段階目で学習した離散化ルールに対応づけるために、Transformerが用いられます。離散化されたあとは、1段階目で学習したデコーダーで情報が復元され、高画質化されます。

ちなみに、このようにベクトル化した画像の離散化（ベクトル量子化）に用いられるルールをコードブック（Codebook）と呼び、そのマッピングにTransformerを用いていることが、CodeFormerの名前の由来となっています。

デモ

それでは、Google Colaboratoryを使って実際に実行していきます。

なお、記事内で紹介したコードをすべて含むノートブックは、以下のリンクから直接参照することができます。

環境設定

はじめに、画面上部のメニューから、「ランタイム」、「ランタイムのタイプを変更」と進み、「ハードウェアアクセラレータ」を「GPU」に変更しておきます。これにより、推論時にGPUを利用することができます。

続いて、公式リポジトリのクローンとライブラリのインストールを行います。

# リポジトリのクローン
!git clone https://github.com/sczhou/CodeFormer.git
%cd CodeFormer

# ライブラリのインストール
!pip install -r requirements.txt
!python basicsr/setup.py develop

モデル設定

学習済みパラメータが公開されているので、ダウンロードします。

!python scripts/download_pretrained_models.py facelib
!python scripts/download_pretrained_models.py CodeFormer

推論

colab上に推論対象のファイルをアップロードします。

続いて、元画像への忠実度CODEFORMER_FIDELITY（高いほど忠実だが低クオリティ）、背景も高画質化するかどうかBACKGROUND_ENHANCEを選択し、アップロードした画像をtest_pathに指定して推論を実行します。

CODEFORMER_FIDELITY = 0.7
BACKGROUND_ENHANCE = True
if BACKGROUND_ENHANCE:
    !python inference_codeformer.py -w $CODEFORMER_FIDELITY -i inputs/user_upload --bg_upsampler realesrgan
else:
    !python inference_codeformer.py -w $CODEFORMER_FIDELITY -i inputs/user_upload

リポジトリ内に用意されているサンプル画像に対する結果がこちらです。