PyCS/pycs/ui/PredictionDialog.py

#!/usr/bin/env python

"""PredictionDialog:  Dialog to select images and videos to predict."""

import os

from PyQt5 import uic, QtWidgets, QtGui, QtCore


class PredictionDialog:
    FileTypeRole = QtCore.Qt.UserRole + 1
    FileNameRole = QtCore.Qt.UserRole + 2

    image_extensions = ['.jpg', '.png', '.jpeg']
    video_extensions = ['.mp4', '.avi', '.mkv', '.mpg', '.mpeg']

    def _get_jobs(self):
        def to_job(item: QtGui.QStandardItem):
            filename = item.data(self.FileNameRole)
            filetype = item.data(self.FileTypeRole)
            return {'filename': filename, 'filetype': filetype, 'jobs': ['detect-faces']}

        return list(map(to_job, [self.model.item(index) for index in range(self.model.rowCount())]))

    jobs = property(fget=_get_jobs)

    def __init__(self, **kwargs):
        self.ui = QtWidgets.QDialog(**kwargs)

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

        self.model = QtGui.QStandardItemModel()
        self.ui.filesList.setModel(self.model)

        # self.add_image('/home/brust/demonstrations/carpe-simiam/images/Alexandra_V110_5_09-08-10.png')
        # self.add_video('/home/brust/demonstrations/carpe-simiam/images/cut.avi')

        self.ui.addFileButton.clicked.connect(self._add_files)

    def _add_files(self):
        extensions = 'All supported files (' + \
                     (' '.join(['*' + extension for extension in self.image_extensions + self.video_extensions])) + \
                     ')'

        selection = QtWidgets.QFileDialog.getOpenFileNames(self.ui, 'Select Files To Add', filter=extensions)
        for filename in selection[0]:
            _, extension = os.path.splitext(filename)
            extension = extension.lower()

            if extension in self.image_extensions:
                self.add_image(filename)
            elif extension in self.video_extensions:
                self.add_video(filename)

    def exec_(self):
        return self.ui.exec_()

    def add_image(self, image):
        # TODO Check for existence
        item = QtGui.QStandardItem(QtGui.QIcon(image), os.path.basename(image))
        item.setData('image', self.FileTypeRole)
        item.setData(image, self.FileNameRole)
        self.model.appendRow(item)

    def add_video(self, video):
        # TODO Check for existence
        item = QtGui.QStandardItem(QtGui.QIcon(video), os.path.basename(video))
        item.setData('video', self.FileTypeRole)
        item.setData(video, self.FileNameRole)
        self.model.appendRow(item)