import json import os from PyQt5 import uic, QtWidgets, QtCore, QtGui from . import AboutDialog from .LabelDialog import LabelDialog from .NewProjectWizard import NewProjectWizard from ..project import Project from ..utils import Video class MainWindow: # Properties def _set_project(self, project): self.__project = project if project is not None: project._update_ui_fn = lambda: self._update_ui() self._update_ui() def _get_project(self): return self.__project _project = property(fget=_get_project, fset=_set_project) def _has_project(self): if self.__project is not None: return True else: return False def _set_predictions(self, predictions): self.__predictions = predictions self._update_prediction_ui() def _get_predictions(self): return self.__predictions _predictions = property(fget=_get_predictions, fset=_set_predictions) def _project_needs_saving(self): return True # Constructor def __init__(self, **kwargs): # initialize window self.ui = QtWidgets.QMainWindow(**kwargs) # set object properties self.__project = None self.__predictions = [] # load ui from corresponding file spath = os.path.dirname(__file__) uic.loadUi(os.path.join(spath, 'MainWindow.ui'), self.ui) # prepare some more ui stuff self.ui.statusLabel = QtWidgets.QLabel(self.ui) self.ui.statusBar.addPermanentWidget(self.ui.statusLabel) self.ui.timerThread = QtCore.QThread(self.ui) self.ui.timer = QtCore.QTimer() self.ui.timer.setInterval(2000) self.ui.timer.moveToThread(self.ui.timerThread) # connect actions # File self.ui.actionNew.triggered.connect(self._project_new) self.ui.actionOpen.triggered.connect(self.project_open) self.ui.actionClose.triggered.connect(self._project_try_closing) self.ui.actionSave.triggered.connect(self._project_save) self.ui.actionQuit.triggered.connect(self._file_quit) # Labeling self.ui.actionStart_Labeling_Session.triggered.connect(self._start_labeling_session) self.ui.startLabelingSessionButton.clicked.connect(self._start_labeling_session) # Help self.ui.actionAbout.triggered.connect(self._help_about) self.ui.actionPredict_Images.triggered.connect(self.predict_via_dialog) self.ui.predictButton.clicked.connect(self.predict_via_dialog) self.ui.imageThumbnailGallery.currentRowChanged.connect(self._prediction_row_changed) self._update_ui() self._update_prediction_ui() # TODO delete me # self._project = Project('test-project') def _update_ui(self): # Status text if self._has_project(): status_text = "OK" if self._project.last_error is not None: status_text = self._project.last_error else: status_text = "No project loaded." self.ui.pipelineStatusLabel.setText(status_text) # Action availability project_only_actions = [self.ui.actionSave, self.ui.actionSave_As, self.ui.actionClose, self.ui.actionImport_Labeled_Faces, self.ui.actionImport_Labeled_Images, self.ui.startLabelingSessionButton, self.ui.actionStart_Labeling_Session, self.ui.actionManage_Individuals, self.ui.actionPredict_Images, self.ui.predictButton, self.ui.actionUpdate_Model, self.ui.updateButton, self.ui.actionValidate_Model, self.ui.clearPredictionsButton, self.ui.exportCSVButton] for action in project_only_actions: action.setEnabled(self._has_project()) def show(self): return self.ui.show() ##################### # Project lifecycle # ##################### def _project_try_closing(self): if self._has_project(): if self._project_needs_saving(): # Ask user box = QtWidgets.QMessageBox() box.setText('The current project has unsaved changes. Would you like to save them?') box.setStandardButtons(QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No | QtWidgets.QMessageBox.Cancel) retval = box.exec_() if retval == QtWidgets.QMessageBox.Yes: self._project.save() if self._project.last_error is not None: error_string = self._project.last_error QtWidgets.QMessageBox.warning(self.ui, 'Error', 'Error while saving (see log for details): %s' % error_string) self._project.reset_error() return False elif retval == QtWidgets.QMessageBox.No: pass elif retval == QtWidgets.QMessageBox.Cancel: return False self._project.close() self._project = None return True else: return True def _project_new(self): retval = self._project_try_closing() if retval: w = NewProjectWizard(self.ui) w_retval = w.exec_() if w_retval: self._project = Project(w.project_root, w.project_config) def project_open(self, path=False): """ opens a project by either using the optional parameter path or or the value the user enters using the ui file dialog :param path: path to project.json """ # check if there is an open project with unsaved changes retval = self._project_try_closing() if not retval: return # copy path to selection or question user if not given if path is not False: selection = path else: selection, _ = QtWidgets.QFileDialog.getOpenFileName(self.ui, 'Select Project Folder', filter='project.json') # check if path exists and open project if len(selection) > 0 and os.path.exists(selection): self._project = Project(selection) if self._project.last_error is not None: error_string = self._project.last_error try: QtWidgets.QMessageBox.warning(self.ui, 'Error', 'Error while opening (see log for details): %s' % error_string) self._project.close() except: pass self._project = None elif len(selection) > 0: # Project file does not exist QtWidgets.QMessageBox.warning(self.ui, 'Error', 'Could not find project file: %s' % selection) def _project_save(self): self._project.save() if self._project.last_error is not None: error_string = self._project.last_error QtWidgets.QMessageBox.warning(self.ui, 'Error', 'Error while saving (see log for details): %s' % error_string) self._project.reset_error() def _file_quit(self): # TODO: warning if there are any unsaved changes self.close() def _start_labeling_session(self): labeling = LabelDialog(self._predictions) labeling.exec_() def _help_about(self): about = AboutDialog() about.exec_() ##################################### # Prediction actions and management # ##################################### def predict_via_dialog(self, selection=False): """ loads images from given paths or ui file dialog :param selection: array of image paths :return: """ # valid image and video extensions image_extensions = ['.jpg', '.png', '.jpeg'] video_extensions = ['.mp4', '.avi', '.mkv', '.mpg', '.mpeg'] # ask user for input if selection is False: name_list = ' '.join(['*' + extension for extension in image_extensions + video_extensions]) extensions = f'All supported files ({name_list})' selection, _ = QtWidgets.QFileDialog.getOpenFileNames(self.ui, 'Select Files To Add', filter=extensions) # load jobs jobs = [] for filename in selection: _, extension = os.path.splitext(filename) extension = extension.lower() if extension in image_extensions: jobs.append({ 'filename': filename, 'filetype': 'image', 'jobs': ['detect-faces'] }) elif extension in video_extensions: jobs.append({ 'filename': filename, 'filetype': 'image', 'jobs': ['detect-faces'] }) progress_dialog = QtWidgets.QProgressDialog('Executing prediction jobs...', 'Abort', 0, 1000, self.ui) progress_dialog.setWindowModality(QtCore.Qt.WindowModal) def callback(progress): progress_dialog.setValue(int(progress * 1000)) for i in range(10): QtWidgets.QApplication.processEvents() self._project.execute(jobs, callback=callback) progress_dialog.close() self._predictions += jobs def _update_prediction_ui(self): # TODO implement better sync self.ui.imageThumbnailGallery.clear() for prediction in self._predictions: filename = prediction['filename'] if prediction['filetype'] == 'image': icon = QtGui.QIcon(filename) item = QtWidgets.QListWidgetItem(icon, os.path.basename(filename)) else: item = QtWidgets.QListWidgetItem(os.path.basename(filename)) self.ui.imageThumbnailGallery.addItem(item) self._prediction_row_changed(self.ui.imageThumbnailGallery.currentRow()) if len(self._predictions) == 0: self._prediction_row_changed(-1) def _prediction_row_changed(self, row): if row in range(len(self._predictions)): prediction = self._predictions[row] self._setup_video_slider(prediction) self._display_prediction(prediction) else: self._setup_video_slider({'filetype': 'empty'}) self._display_prediction({'filetype': 'empty'}) def _setup_video_slider(self, prediction): if prediction['filetype'] == 'video': video_cap = Video(prediction['filename']) if video_cap.last_error is None: self.ui.frameNumber.setDecMode() self.ui.frameNumber.display(0) self.ui.frameNumber.setDigitCount(5) self.ui.frameNumber.show() self.ui.frameSlider.show() self.ui.frameSlider.setValue(0) self.ui.frameSlider.setMinimum(0) self.ui.frameSlider.setMaximum(video_cap.framecount - 1) self.ui.frameSlider.valueChanged.connect( lambda value, _prediction=prediction: [self._display_prediction(dict(frame=value, cap=video_cap, **_prediction)), self.ui.frameNumber.display(value)]) else: self._disable_video_slider() else: self._disable_video_slider() def _disable_video_slider(self): self.ui.frameSlider.setValue(0) try: self.ui.frameSlider.valueChanged.disconnect() except: pass self.ui.frameSlider.hide() self.ui.frameNumber.hide() def _display_prediction(self, prediction): self.ui.annotatedImageView.display(prediction) filtered_dict = {i: prediction[i] for i in prediction if i != 'cap'} if 'cap' in prediction.keys(): filtered_dict['cap'] = '[VideoCapture]' self.ui.predictionDebugText.setPlainText(json.dumps(filtered_dict, indent=2, skipkeys=True))