PyCS/pycs/ui/MainWindow.py

#!/usr/bin/env python

import copy
import json
import os

from PyQt5 import uic, QtWidgets, QtCore, QtGui

from . import AboutDialog
from .NewProjectWizard import NewProjectWizard
from .PredictionDialog import PredictionDialog
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):
        self.ui = QtWidgets.QMainWindow(**kwargs)

        # Properties
        self.__project = None
        self.__predictions = []

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

        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)

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

        # 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):
        retval = self._project_try_closing()
        if retval:
            selection = QtWidgets.QFileDialog.getExistingDirectory(self.ui, 'Select Project Folder')
            if len(selection) > 0 and os.path.exists(os.path.join(selection, 'project.json')):
                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' % os.path.join(selection, 'project.json'))

    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 _help_about(self):
        about = AboutDialog()
        about.exec_()

    #####################################
    # Prediction actions and management #
    #####################################
    def _predict_via_dialog(self):
        prediction_dialog = PredictionDialog(parent=self.ui)
        if prediction_dialog.exec_():

            # TODO extract callback and progress dialog to make it reusable
            jobs = copy.deepcopy(prediction_dialog.jobs)
            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))