|
@@ -1,335 +0,0 @@
|
|
|
-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))
|