MXNetを使ってMSRNで超解像

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

github.com

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

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

実行スクリプト

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

ctx = mx.cpu()

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('MSRN_4x', 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
prob = (prob.transpose(1,2,0)*255).astype(np.uint8)
img = Image.fromarray(prob)
img.save('MSRN_4x.jpg')

結果の表示

いままでで一番きれいなような気がする。