MXNet-Gluon-Style-TransferをGUIで動かす

初めに

MXNet-Gluon-Style-Transferについてはこちらを参照
touch-sp.hatenablog.com

環境

Windows10 Pro 64bit (GPUなし)
Python 3.6.8

バージョン確認(pip freeze)

certifi==2019.6.16
chardet==3.0.4
graphviz==0.8.4
idna==2.6
mxnet==1.4.1
numpy==1.16.4
Pillow==6.0.0
PyQt5==5.12.2
PyQt5-sip==4.19.17
requests==2.18.4
urllib3==1.22

「qt.py」

import sys
from PyQt5 import QtCore
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QImage, QPixmap

import main_eval

class App(QWidget):

    def __init__(self):
        super().__init__()
        self.title = "PyQt5 sample"
        self.left = 50
        self.top = 50
        self.width = 1200
        self.height = 660
        self.initUI()
        self.style_img = ''
        self.content_img = ''

    def initUI(self):
        self.setWindowTitle(self.title)
        self.setGeometry(self.left, self.top, self.width, self.height)
        
        self.button1 = QPushButton('Open Img', self)
        self.button1.setGeometry(QtCore.QRect(10, 10, 97, 27))

        self.button2 = QPushButton('Style Select', self)
        self.button2.setGeometry(QtCore.QRect(130, 10, 97, 27))
        
        self.button3 = QPushButton('style-transfer', self)
        self.button3.setGeometry(QtCore.QRect(560, 360, 81, 27))
        
        self.img_label1 = QLabel(self)
        self.img_label1.setGeometry(QtCore.QRect(20, 120, 512, 512))
        self.img_label1.setFrameStyle(QFrame.Box | QFrame.Plain)

        self.img_label2 = QLabel(self)
        self.img_label2.setGeometry(QtCore.QRect(660, 120, 512, 512))
        self.img_label2.setFrameStyle(QFrame.Box | QFrame.Plain)

        self.button1.clicked.connect(self.showDialog1)
        self.button2.clicked.connect(self.showDialog2)
        self.button3.clicked.connect(self.style_transfer)

        self.show()       

    def showDialog1(self):

        fname = QFileDialog.getOpenFileName(self, 'Open file')

        # fname[0]は選択したファイルのパス(ファイル名を含む)
        if fname[0]:
            self.style_img= fname[0]
            # 画像の読み込み, サイズ変更
            image = QImage(fname[0]).scaled(512,512,QtCore.Qt.KeepAspectRatio) 
            # 画像の表示
            self.img_label1.setPixmap(QPixmap.fromImage(image))

    def showDialog2(self):

        fname = QFileDialog.getOpenFileName(self, 'Open file')

        # fname[0]は選択したファイルのパス(ファイル名を含む)
        if fname[0]:
            self.content_img=fname[0]
            # 画像の読み込み, サイズ変更
            image = QImage(fname[0]).scaled(512,512,QtCore.Qt.KeepAspectRatio) 
            # 画像の表示
            self.img_label2.setPixmap(QPixmap.fromImage(image))
    
    def style_transfer(self):
        main_eval.evaluate(self.style_img, self.content_img)
       
if __name__ == "__main__":
    app = QApplication(sys.argv)
    ex = App()
    sys.exit(app.exec_())

「main_eval.py」

import net
import utils
import mxnet as mx

def evaluate(content, style):
    
    ctx = mx.cpu(0)
    
    # images
    content_image = utils.tensor_load_rgbimage(content,ctx, size=512, keep_asp=True)
    style_image = utils.tensor_load_rgbimage(style, ctx, size=512)
    style_image = utils.preprocess_batch(style_image)
    # model
    style_model = net.Net(ngf=128)
    style_model.load_parameters('models/21styles.params', ctx=ctx)
    # forward
    style_model.set_target(style_image)
    output = style_model(content_image)
    utils.tensor_save_bgrimage(output[0], 'output.jpg', 0)