PyCS/pycs/ui/LabelDialog.py

import os

from PyQt5 import uic, QtWidgets, QtCore

from pycs.ui.AnnotatedInteractiveImageView import AnnotatedInteractiveImageView


class LabelDialog(QtWidgets.QDialog):
    def __init__(self, predictions, **kwargs):
        # call parent constructor
        super(LabelDialog, self).__init__(**kwargs)

        # handle properties
        self.__predictions = predictions
        self.__current = -1

        # load ui
        spath = os.path.dirname(__file__)
        uic.loadUi(os.path.join(spath, 'LabelDialog.ui'), self)

        self.button_next.clicked.connect(self._next)
        self.button_previous.clicked.connect(self._prev)

        # add annotated image view
        self.__image = AnnotatedInteractiveImageView(self, self.click)
        self.gridLayout.addWidget(self.__image)

        # call next image function to display first item
        self._next()

    def _next(self):
        """
        displays next item
        """
        if len(self.__predictions) > 0:
            self.__current += 1
            if self.__current >= len(self.__predictions):
                self.__current = 0

            self.__image.display(self.__predictions[self.__current])

    def _prev(self):
        """
        displays previous item
        """
        if len(self.__predictions) > 0:
            self.__current -= 1
            if self.__current <= 0:
                self.__current = len(self.__predictions) - 1

            self.__image.display(self.__predictions[self.__current])

    def click(self, x, y, event):
        """
        handles clicks on images
        :param x: x position in [0, 1]
        :param y: y position in [0, 1]
        :param event: actual QMouseEvent object
        """

        # return if there is no prediction
        if len(self.__predictions) == 0:
            return

        # add label object to current prediction
        if 'label' not in self.__predictions[self.__current].keys():
            self.__predictions[self.__current]['label'] = {
                'faces': [{
                    'x1': None,
                    'y1': None,
                    'x2': None,
                    'y2': None
                }]
            }

        face = self.__predictions[self.__current]['label']['faces'][0]

        # set data depending on click
        if event.button() == QtCore.Qt.LeftButton:
            face['x1'] = x
            face['y1'] = y
        elif event.button() == QtCore.Qt.RightButton:
            face['x2'] = x
            face['y2'] = y

        # return if not at least two points selected
        if face['x1'] is None or face['y1'] is None or face['x2'] is None or face['y2'] is None:
            return

        # detect upper left point and calculate width and height
        face['x'] = min(face['x1'], face['x2'])
        face['y'] = min(face['y1'], face['y2'])
        face['w'] = max(face['x1'], face['x2']) - face['x']
        face['h'] = max(face['y1'], face['y2']) - face['y']

        # display rectangle
        self.__image.display(self.__predictions[self.__current], mode='label')