MXNetを使ってMDSRで超解像

はじめに

• EDSR、MDSRについてはこちらを参照

[1707.02921] Enhanced Deep Residual Networks for Single Image Super-Resolution

• MXNetの学習済みモデルはこちらからダウンロード可能

github.com

サンプル画像のダウンロード

『こちら』からダウンロードして「dog.jpg」の名前で保存

実行スクリプト

import numpy as np
import mxnet as mx
from mxnet import image

ctx = mx.cpu()
#model_name = 'MDSR_x2'
model_name = 'MDSR_x4'

img = image.imread('dog.jpg')
img = img.astype(np.float32)/255
img = mx.nd.transpose(img, (2,0,1))
img = mx.nd.expand_dims(img, axis=0)

sym, arg_params, aux_params = mx.model.load_checkpoint(model_name, 0)
model = mx.mod.Module(symbol=sym, label_names=None, context=ctx)
model.bind(for_training=False, data_shapes=[('data', img.shape)])
model.set_params(arg_params, aux_params)

from collections import namedtuple
Batch = namedtuple('Batch', ['data'])

model.forward(Batch([img]), is_train=False)
prob = model.get_outputs()[0].asnumpy()
prob = np.squeeze(prob)

from PIL import Image
output_file = model_name + '.jpg'
prob = (prob.transpose(1,2,0)*255).astype(np.uint8)
img = Image.fromarray(prob)
img.save(output_file)

結果の表示

• 目の部分だけ拡大表示
• 左から「元画像」「2倍の超解像」「4倍の超解像」

from PIL import Image

original = Image.open('dog.jpg').resize((1024,1024))
x0 = original.crop((580,360,700,500))

x2_original = Image.open('MDSR_x2.jpg').resize((1024,1024))
x2 = x2_original.crop((580,360,700,500))

x4_original = Image.open('MDSR_x4.jpg')
x4 = x4_original.crop((580,360,700,500))

img = Image.new('RGB', (120*3, 140))
img.paste(x0, (0, 0))
img.paste(x2, (120, 0))
img.paste(x4, (240, 0))

img.save('result.jpg')