color_transfer_ui.py

#!/usr/bin/python

"""
color transfer ui
"""

import os
import sys
import cv2
import tempfile
from color_transfer import color_transfer as ct
from PyQt4.QtCore import *
from PyQt4.QtGui import *
reload(sys)
sys.setdefaultencoding('utf-8')

__version__ = '1.0.0'

class MainWindow(QMainWindow):
    targetFile = ""
    sourceFile = ""

    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)

        # set main window size
        desktop = QDesktopWidget()
        self.dwidth = desktop.width()
        self.dheight = desktop.height()
        self.resize(self.dwidth, self.dheight)

        sourceFileOpenAction = self.createAction("Source", slot=self.sourceFileOpen, shortcut="Ctrl+1", icon=None, tip="Open an image file")
        targetFileOpenAction = self.createAction("Target", slot=self.targetFileOpen, shortcut="Ctrl+2", icon=None, tip="Open an image file") 
        transferImgAction = self.createAction("Transfer", slot=self.transferImg, shortcut="Ctrl+3", icon=None, tip="Transfer source image color to target image") 

        # menubar
        fileMenu = self.menuBar().addMenu("&File")
        fileMenu.addAction(sourceFileOpenAction)
        fileMenu.addAction(targetFileOpenAction)
        fileMenu.addAction(transferImgAction)
        #self.addActions(fileMenu, (sourceFileOpenAction, targetFileOpenAction))
        
        # toolbar
        fileToolbar = self.addToolBar("File")
        fileToolbar.setObjectName("FileToolBar")
        fileToolbar.addAction(sourceFileOpenAction)
        fileToolbar.addAction(targetFileOpenAction)
        fileToolbar.addAction(transferImgAction)

        gridLayout = QGridLayout()

        self.sourceLabel = QLabel()
        self.targetLabel = QLabel()
        self.transferLabel = QLabel()

        gridLayout.addWidget(self.sourceLabel, 0, 0)
        gridLayout.addWidget(self.targetLabel, 0, 1)
        gridLayout.addWidget(self.transferLabel, 0, 2)

        mainWidget = QWidget()
        mainWidget.setLayout(gridLayout)
        self.setCentralWidget(mainWidget)


    def addActions(self):
        pass

    def transferImg(self):
        source = cv2.imread(self.sourceFile)
        target = cv2.imread(self.targetFile)
        transfer = ct(source, target)
        outpic = tempfile.gettempdir() + os.path.sep + 'tmp.jpg'
        #outpic = os.path.abspath(__file__)
        cv2.imwrite(outpic, transfer)
        self.showImage(outpic, self.transferLabel)

    def sourceFileOpen(self):
        self.sourceFile = self.fileOpen()
        self.showImage(self.sourceFile, self.sourceLabel)

    def targetFileOpen(self):
        self.targetFile = self.fileOpen()
        self.showImage(self.targetFile, self.targetLabel)

    def fileOpen(self):
        #dir = os.path.dirname(self.filename) if self.filename is not None else '.'
        dir = os.path.dirname(__file__) 
        formats = ["*.%s" % unicode(format).lower() for format in QImageReader.supportedImageFormats()]
        fname = unicode(QFileDialog.getOpenFileName(self, "Image - Choose Image", dir, "Images (*.png *.jpg)"))
        return fname

    def showImage(self, fname, imgLoader):
        image = QImage(fname)
        if image.isNull():
            print 'error to load image'
        else:
            self.image = QImage()
            self.image = image
            self.image.filename = fname
            percent = 100
            factor = percent / 100.0
            width = self.image.width() * factor
            height = self.image.height() * factor
            image = self.image.scaled(self.dwidth / 3.0 ,self.dheight, Qt.KeepAspectRatio)
            #image = self.image.scaled(width, height, Qt.KeepAspectRatio)
            imgLoader.setMaximumSize(self.dwidth / 3.0, self.dheight)
            imgLoader.setMinimumSize(200, 200)
            imgLoader.setAlignment(Qt.AlignCenter)
            imgLoader.setContextMenuPolicy(Qt.ActionsContextMenu)
            imgLoader.setPixmap(QPixmap.fromImage(image))
            #self.imageLabel.setPixmap(QPixmap.fromImage(image))


    def createAction(self, text, slot=None, shortcut=None, icon=None, tip=None, checkable=False, signal="triggered()"):
        action = QAction(text, self)
        if icon is not None:
            action.setIcon(QIcon(":/%s.png") % icon)
        if shortcut is not None:
            action.setShortcut(shortcut)

        if tip is not None:
            action.setToolTip(tip)
            action.setStatusTip(tip)
        if slot is not None:
            self.connect(action, SIGNAL(signal), slot)
        action.setCheckable(checkable)
        return action

def main():
    app = QApplication(sys.argv)
    form = MainWindow()
    form.show()
    app.exec_()

if __name__ == '__main__':
    main()