公開日:2021年11月17日
最終更新日:2022年9月15日
はじめに
以前MXNetを使った超解像の記事を書きました。touch-sp.hatenablog.com
touch-sp.hatenablog.com
今回はPyTorchとMMEditingを使ってみたいと思います。
環境
二つの環境で動作確認できています。Ubuntu on WSL2
Ubuntu 20.04 LTS on WSL2 (Windows 11) CUDA 11.3.1
Python 3.9.5
Windows
Windows 11 CUDA 11.6.2
Python 3.9.13
Python環境構築
すべてpipでインストール可能です。(CUDAのバージョンを自身の環境に合わせて下さい)
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113 pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.12.0/index.html pip install mmedit pip install openmim pip install mmengine
Pythonスクリプト
Python環境が整えばあとは以下のスクリプトを実行するのみです。リポジトリのクローン(git clone)などは必要ありません。
必要なファイルはスクリプト内でダウンロードできるようになっています。
画像の部分は好きな画像に変更してください。
import os from PIL import Image import torch from torchvision.datasets.utils import download_url from mmedit.apis import init_model, restoration_inference from mmedit.core import tensor2img from mim.commands.download import download device = 'cuda' if torch.cuda.is_available() else 'cpu' img_url ='https://github.com/open-mmlab/mmediting/raw/master/tests/data/lq/baboon_x4.png' img_fname = img_url.split('/')[-1] download_url(img_url, root = '.', filename = img_fname) os.makedirs('models', exist_ok=True) checkpoint_name = 'esrgan_x4c64b23g32_g1_400k_div2k' config_fname = checkpoint_name + '.py' checkpoint = download(package="mmedit", configs=[checkpoint_name], dest_root="models")[0] model = init_model(os.path.join('models', config_fname), os.path.join('models', checkpoint), device = device) output = restoration_inference(model, img_fname) output = tensor2img(output) pil_img = Image.fromarray(output[:,:,::-1]) pil_img.show()