`handbag_data/github_datasets/iGAN/ui/gui_vis.py`

114 lines
import numpy as np
import cv2
from PyQt4.QtCore import *
from PyQt4.QtGui import *
from . import save_result
from lib import utils


class GUI_VIS(QWidget):
    def __init__(self, opt_engine, topK=16, grid_size=None, nps=320, model_name='tmp'):
        QWidget.__init__(self)
        self.topK = topK
        if grid_size is None:
            self.n_grid = int(np.ceil(np.sqrt(self.topK)))
            self.grid_size = (self.n_grid, self.n_grid) # (width, height)
        else:
            self.grid_size = grid_size
        self.select_id = 0
        self.ims = None
        self.vis_results = None
        self.width = int(np.round(nps/ (4 * float(self.grid_size[1])))) * 4
        self.winWidth = self.width * self.grid_size[0]
        self.winHeight = self.width * self.grid_size[1]

        self.setFixedSize(self.winWidth, self.winHeight)
        self.opt_engine = opt_engine
        self.frame_id = -1
        self.sr = save_result.SaveResult(model_name=model_name)

    def save(self):
        self.sr.cache_result(ims=self.ims, visIms=self.vis_results)
        self.sr.save()

    def set_frame_id(self, frame_id):
        if frame_id != self.frame_id:
            self.frame_id = frame_id
            self.update_vis()

    def set_image_id(self, image_id):
        if image_id != self.select_id:
            self.select_id = image_id
            self.update_vis()

    def paintEvent(self, event):
        painter = QPainter()
        painter.begin(self)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.fillRect(event.rect(), Qt.white)
        if self.vis_results is not None:
            qImg = QImage(self.vis_results.tostring(), self.winWidth, self.winHeight, QImage.Format_RGB888)
            painter.drawImage(0, 0, qImg)

        painter.end()

    def update_results(self, ims):
        self.ims = ims

    def mousePressEvent(self, event):
        pos = event.pos()
        if event.button() == Qt.LeftButton:
            x_select = np.floor(pos.x() / float(self.width))
            y_select = np.floor(pos.y() / float(self.width))
            new_id = int(y_select * self.grid_size[0] + x_select)
            print('pos=(%d,%d) (x,y)=(%d,%d) image_id=%d' % (int(pos.x()), int(pos.y()), x_select, y_select, new_id))
            if new_id != self.select_id:
                self.select_id = new_id
                self.update_vis()
                self.update()
                self.emit(SIGNAL('update_image_id'), self.select_id)

    def sizeHint(self):
        return QSize(self.winWidth, self.winHeight)

    def reset(self):
        self.ims = None
        self.vis_results = None
        self.update()
        self.sr.reset()
        self.frame_id = -1
        self.select_id=0

    def get_show_id(self):
        return self.select_id

    def set_show_id(self, _id):
        if _id != self.select_id:
            self.select_id = _id
            self.show_results()

    def update_vis(self):
        ims = self.opt_engine.get_images(self.frame_id)

        if ims is not None:
            self.ims = ims

        if self.ims is None:
            return

        ims_show = []
        n_imgs = self.ims.shape[0]
        for n in range(n_imgs):
            # im = ims[n]
            im_s = cv2.resize(self.ims[n], (self.width, self.width), interpolation=cv2.INTER_CUBIC)
            if n == self.select_id and self.topK > 1:
                t = 3  # thickness
                cv2.rectangle(im_s, (t, t), (self.width - t, self.width - t), (0, 255, 0), t)
            im_s = im_s[np.newaxis, ...]
            ims_show.append(im_s)
        if ims_show:
            ims_show = np.concatenate(ims_show, axis=0)
            g_tmp = utils.grid_vis(ims_show, self.grid_size[1], self.grid_size[0]) # (nh, nw)
            self.vis_results = g_tmp.copy()
            self.update()