今回は高解像度化ができるRealESRGANをgoogle colab上で試してみました
この記事はRealESRGANを動作させるまでの備忘録になります

この記事でわかること

google colab上でRealESRGANを使って画像を高画質化する方法

1.実行環境

Google Colab
git:clone https://github.com/xinntao/Real-ESRGAN.git
model : RealESRGAN_x4plus

2.環境構築＆コード記述

2.1 インストール＆ダウンロード

以下コマンドでgitからcloneして必要モジュールをインストールします
またモデルで使用する重みもダウンロードします

!git clone https://github.com/xinntao/Real-ESRGAN.git
%cd Real-ESRGAN

!apt-get install ffmpeg
!pip install pydub
!pip install basicsr
!pip install facexlib
!pip install gfpgan
!pip install -r requirements.txt
!python setup.py develop

!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models

2.2 モデル実行～結果保存

指定モデルや引数を読み込んで、高画質化して保存する処理を記述します

import cv2
import glob
import os
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
import subprocess
from pydub import AudioSegment
from google.colab.patches import cv2_imshow

model_path = "experiments/pretrained_models/RealESRGAN_x4plus.pth" #@param {type:"string"}
path = 'path/to/image/'  #@param {type:"string"}

# variable
block = 23
netscale = 4
tile = 0
tile_pad = 10
pre_pad = 0
half = False
face_enhance = False
outscale = 8 #@param {type:"number"}
output_dir = '/path/to/output_dir' #@param {type:"string"}
resize_rate = 0.5
limit_img_size = 1000

if 'RealESRGAN_x4plus_anime_6B.pth' in model_path:
    block = 6
elif 'RealESRGAN_x2plus.pth' in model_path:
    netscale = 2

model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=block, num_grow_ch=32, scale=netscale)

upsampler = RealESRGANer(
    scale=netscale,
    model_path=model_path,
    model=model,
    tile=tile,
    tile_pad=tile_pad,
    pre_pad=pre_pad,
    half=half)

if face_enhance:
    from gfpgan import GFPGANer
    face_enhancer = GFPGANer(
        model_path='path/to/model',
        upscale=outscale,
        arch='clean',
        channel_multiplier=2,
        bg_upsampler=upsampler)
else:
  pass    

os.makedirs(output_dir, exist_ok=True)

img = cv2.imread(path)
height = img.shape[0]
width = img.shape[1]
print('original image:', height, width)
if height < width:
  long_size = width
else:
  long_size = height

# 入力画像のサイズが1000pixel以上だとメモリオーバーになるのでリサイズする  
if limit_img_size < long_size:
  resize_rate = limit_img_size/long_size
  in_img = cv2.resize(img , (int(width*resize_rate), int(height*resize_rate)))
  print('resize image', in_img.shape[0], in_img.shape[1])
else:
  in_img = img

if face_enhance:
    # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
    pass
else:
    output, _ = upsampler.enhance(in_img, outscale=outscale)

save_path = path/to/save
cv2.imwrite(save_path, output)
del output
del in_img
del img
del upsampler