本記事では、DDNMと呼ばれる機械学習手法を用いて画像に超解像処理を施す方法をご紹介します。

出典: wyhuai/DDNM

DDNM

概要

Denoising Diffusion Null-Space Model (DDNM)は、様々な画像復元タスクを実現するゼロショットフレームワークです。

出典: wyhuai/DDNM

従来技術は超解像やカラー化などそれぞれのタスクで固有のモデルを構成しており、様々なタスクを一般化することはできませんでした。

DDNMでは、超解像(SuperResolution)、カラー化(colorization)、修復(Inpainting)、圧縮センシング(compressed sensing)、ブレ除去(deblurring)といった様々な画像復元タスクを実現します。

DDNMでは、追加のトレーニングやネットワークの変更を必要とせず、事前にトレーニングされた一般の拡散モデルのgenerative priorのみを用います。逆拡散過程でヌル空間の内容を精製することで、データの一貫性とリアリティを両立させています。

出典: Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model

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

本記事では上記手法を用いて、画像の超解像を行っていきます。

デモ(Colaboratory)

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

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

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

環境セットアップ

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

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

%cd /content

!git clone https://github.com/wyhuai/DDNM.git

# Commits on Dec 11, 2022
%cd /content/DDNM
!git checkout f157d760710169c9ba22a830caac9035ade71b5f

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

%cd /content/DDNM

!pip install blobfile==2.0.0 tqdm==4.61.1 pyYaml==6.0 pillow==7.1.2

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

import glob
from PIL import Image
from matplotlib import pyplot as plt
import numpy as np

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

学習済みモデルのセットアップ

続いて、論文発表元が公開する学習済みモデルをGoogle Colaboratoryにダウンロードします。

%cd /content/DDNM

# Create dir
!mkdir -p DDNM/exp/logs/celeba DDNM/exp/logs/imagenet

# Download model
!wget -c https://image-editing-test-12345.s3-us-west-2.amazonaws.com/checkpoints/celeba_hq.ckpt \
      -O DDNM/exp/logs/celeba/celeba_hq.ckpt
!wget -c https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt \
      -O DDNM/exp/logs/imagenet/256x256_diffusion_uncond.pt

Super Resolution

それでは、超解像を行っていきます。
celebaデータセットでトレーニングされたモデルを使用し、人の顔の超解像(x4)を行っていきます。

以下では、./exp/datasets/celeba_hqに格納されたファイルに超解像を行います。

# base path
exp = './exp'
# input dir {exp}/datasets/{in_dir}
in_dir = 'celeba_hq'
# output dir
out_dir = 'sr_results' # {exp}/image_sample/{out_dir}
# tasks
tasks = 'sr_averagepooling' #@param ['cs_walshhadamard', 'cs_blockbased', 'inpainting', 'denoising', 'deblur_uni', 'deblur_gauss', 'deblur_aniso', 'sr_averagepooling', 'sr_bicubic', 'colorization', 'mask_color_sr']
# config path configs/{config_path}
config_path = 'celeba_hq.yml'

%cd /content/DDNM

!python main.py \
  --ni --simplified \
  --config {config_path} \
  --path_y {in_dir} \
  --eta 0.85 \
  --deg {tasks} \
  --deg_scale 4.0 --sigma_y 0 \
  --exp {exp} \
  -i {out_dir}

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

input_imgs = glob.glob(f'{exp}/image_samples/{out_dir}/Apy/Apy_*.png')
input_imgs.sort()
output_imgs = glob.glob(f'{exp}/image_samples/{out_dir}/*.png')
output_imgs.sort()

for i, o in zip(input_imgs, output_imgs):
  in_img = Image.open(i)
  out_img = Image.open(o)

  fig = plt.figure(num=None, figsize=(12, 5))

  ax = fig.add_subplot(1, 2, 1, xticks=[], yticks=[])
  plt.imshow(np.asarray(in_img))
  ax.set_title('input')

  ax = fig.add_subplot(1, 2, 2, xticks=[], yticks=[])
  plt.imshow(np.asarray(out_img))
  ax.set_title('4x output')

機械的に低画質化された入力画像を超解像できていることが確認できます。

Old Photo Restoration

続いて、古い写真の復元を動かしていきます。

# base path
exp = './exp'
# input dir {exp}/datasets/{in_dir}
in_dir = 'oldphoto'
# output dir
out_dir = 'opr_results' # {exp}/image_sample/{out_dir}
# tasks
tasks = 'mask_color_sr' #@param ['cs_walshhadamard', 'cs_blockbased', 'inpainting', 'denoising', 'deblur_uni', 'deblur_gauss', 'deblur_aniso', 'sr_averagepooling', 'sr_bicubic', 'colorization', 'mask_color_sr']
# config path configs/{config_path}
config_path = 'oldphoto.yml'

%cd /content/DDNM

!python main.py \
  --ni --simplified \
  --config {config_path} \
  --path_y {in_dir} \
  --eta 0.85 \
  --deg {tasks} \
  --deg_scale 2.0 --sigma_y 0.02 \
  --exp {exp} \
  -i {out_dir}

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

input_imgs = glob.glob(f'{exp}/image_samples/{out_dir}/Apy/Apy_*.png')
input_imgs.sort()
output_imgs = glob.glob(f'{exp}/image_samples/{out_dir}/*.png')
output_imgs.sort()

for i, o in zip(input_imgs, output_imgs):
  in_img = Image.open(i)
  out_img = Image.open(o)

  fig = plt.figure(num=None, figsize=(12, 5))

  ax = fig.add_subplot(1, 2, 1, xticks=[], yticks=[])
  plt.imshow(np.asarray(in_img))
  ax.set_title('input')

  ax = fig.add_subplot(1, 2, 2, xticks=[], yticks=[])
  plt.imshow(np.asarray(out_img))
  ax.set_title('old photo restoration output')

モデルはそのままに、異なるタスクが実行できることが確認できます。
ただし、実行速度や復元精度は専用タスクである以下の手法の方が使い勝手が良いかもしれません。

まとめ

本記事では、DDNMを用いて超解像や、古い写真の復元を行う方法をご紹介しました。
推論速度や、精度の改善により高い汎用性が期待できます。

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

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

参考文献

1.  論文 - Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model

2. GitHub - wyhuai/DDNM