[Python] Real-ESRGANを使ってみる – 野田ファームの備忘録

高解像度化（アップスケーリング）を行うAIモデル「Real-ESRGAN」をローカル環境で使ってみました。

※NVIDIA GeForce RTX 3060を使用しています。

Python 3.10.11をインストール

Real-ESRGANを使う場合、Pythonは3.8以上を推奨しているので今回は3.10.11をインストールします。

https://www.python.org/downloads/windows

インストール出来たらバージョンを確認します。

python --version

CUDA11.8をインストール

https://developer.nvidia.com/cuda-11-8-0-download-archive

インストール出来たらバージョンを確認します。

nvcc --version

プロジェクトフォルダを作成

「real-esrgan」というフォルダを作成します。

mkdir real-esrgan

「real-esrgan」フォルダに移動します。

cd real-esrgan

スクリプトを配置するためのフォルダ「src」を作成します。

mkdir src

事前学習済みモデルを配置するためのフォルダ「weights」を作成します。

mkdir weights

入力用画像を配置するためのフォルダ「input_images」を作成します。

mkdir input_images

出力用画像を配置するためのフォルダ「output_images」を作成します。

mkdir output_images

仮想環境を作成

「venv」という名前で仮想環境を作成します。

python -m venv venv

完了すると「real-esrgan」フォルダ内に「venv」という名前のフォルダが作成されます。

仮想環境を有効化

「venv」フォルダ内にある「Scripts」フォルダに移動します。

cd venv/scripts

仮想環境を有効化します。

activate

有効化されると「(venv) ・・・」と表示されます。

「venv」フォルダに戻ります。

cd..

ライブラリをインストール

まずはpipを最新版にアップデートします。

python.exe -m pip install --upgrade pip

CUDA 11.8対応版の必要なライブラリをインストールします。

pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118

Real-ESRGANをインストール

GitHubからクローンします。

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

「venv」フォルダ内に「Real-ESRGAN」フォルダがあるので移動します。

cd real-esrgan

Real-ESRGANに依存するライブラリをインストールします。

pip install -r requirements.txt

Real-ESRGANをpipでインストールします。

pip insatall -e .

Real-ESRGANがインストールされたことを確認します。

pip list

リストの中に「realesrgan」があればOKです。

プロジェクトフォルダまで戻ります。

cd../..

CUDAが利用可能か確認

「src」フォルダ内に「check.py」ファイルを作成して、以下のコードを記述します。

import torch
import torchvision

print("Torch version:", torch.__version__)
print("Torchvision version:", torchvision.__version__)
print("CUDA available:", torch.cuda.is_available())

「check.py」を実行します。

python src/check.py

以下のような結果が出ればOKです。

Torch version: 2.0.1+cu118
Torchvision version: 0.15.2+cu118
CUDA available: True

事前学習済みモデルを用意

今回はアニメや漫画のようなイラストをアップスケーリングしてみたいので、それに特化したモデルを使います。

https://github.com/xinntao/Real-ESRGAN/releases

V0.2.2.4にある「RealESRGAN_x4plus_anime_6B.pth」をダウンロードします。

ダウンロードしたファイルを「weights」フォルダ内に置きます。

入力用画像を用意

今回はChatGPTで画像を生成したものを用意しました。

画像サイズは600×600にしました。

画像ファイル名を「test.jpg」で、「input_images」フォルダ内に格納します。

アップスケーリング

「src」フォルダ内に「main.py」ファイルを作成して、以下のコードを記述します。

import torch
from PIL import Image
import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

# モデルの初期化
device = 'cuda' if torch.cuda.is_available() else 'cpu'

model_path = 'weights/RealESRGAN_x4plus_anime_6B.pth'
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)

upscaler = RealESRGANer(
    scale=4,
    model_path=model_path,
    dni_weight=None,
    model=model,
    tile=256,  # タイル処理を有効化
    tile_pad=10,
    pre_pad=0,
    half=True if device == 'cuda' else False,
    device=device
)

# 入力画像を読み込む
input_image = Image.open('input_images/test.jpg').convert('RGB')

# NumPy配列に変換
input_image_array = np.array(input_image)

# 高解像度化処理
try:
    output_image_array, _ = upscaler.enhance(input_image_array)
    # 結果を保存
    output_image = Image.fromarray(output_image_array)
    output_image.save('output_images/test.jpg')
    print("高解像度化が完了しました: test.jpg")
except Exception as e:
    print(f"エラーが発生しました: {e}")

「main.py」を実行します。

python src/main.py

「高解像度化が完了しました: test.jpg」と表示されたら正常に処理されています。

「output_images」フォルダ内に「test.jpg」ファイルがあるので確認してみましょう。

約3秒ほどで4倍のサイズにアップスケーリングされました。

画像編集ソフトで拡大した場合とは違い、かなりきれいにアップスケーリングされました。

おまけ1

Real-ESRGANの汎用モデルを使いたい場合は「RealESRGAN_x4plus.pth」になります。

num_blockの値は「23」に変えます。

model_path = 'weights/RealESRGAN_x4plus.pth'
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)

おまけ2

「input_images」フォルダ内にある画像ファイルをすべてアップスケーリングするようにしてみました。

import os
import torch
from PIL import Image
import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

# モデルの初期化
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 入力と出力のフォルダ
input_folder = 'input_images'
output_folder = 'output_images'
os.makedirs(output_folder, exist_ok=True)

model_path = 'weights/RealESRGAN_x4plus_anime_6B.pth'
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)

upscaler = RealESRGANer(
    scale=4,
    model_path=model_path,
    dni_weight=None,
    model=model,
    tile=256,  # タイル処理を有効化
    tile_pad=10,
    pre_pad=0,
    half=True if device == 'cuda' else False,
    device=device
)

for file_name in os.listdir(input_folder):
    if file_name.endswith(('.png', '.jpg', '.jpeg')):
        input_path = os.path.join(input_folder, file_name)
        output_path = os.path.join(output_folder, file_name)

        # 入力画像を読み込む
        input_image = Image.open(input_path).convert('RGB')

        # NumPy配列に変換
        input_image_array = np.array(input_image)

        # 高解像度化処理
        try:
            output_image_array, _ = upscaler.enhance(input_image_array)
            # 結果を保存
            output_image = Image.fromarray(output_image_array)
            output_image.save(output_path)
        except Exception as e:
            print(f"エラーが発生しました: {e}")

おまけ3

仮想環境を終了するときは以下を実行します。

cd venv/scripts
deactivate