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')