PyCS/pycs/ui/AnnotatedImageView.py

"""AnnotatedImageView:  Displays images with bounding boxes etc."""

from PyQt5 import QtWidgets, QtGui, QtCore

from ..utils import Video


class AnnotatedImageView(QtWidgets.QGraphicsView):
    def __init__(self, parent):
        super(AnnotatedImageView, self).__init__(parent=parent)

        self._scene = QtWidgets.QGraphicsScene()
        self.setScene(self._scene)

        self._rect = QtCore.QRectF(0, 0, 0, 0)

        self.setRenderHint(QtGui.QPainter.Antialiasing)

        self._image = None
        self._current_video = None

    def display(self, prediction, mode='prediction'):
        self._scene.clear()
        if prediction['filetype'] == 'image':
            self._image = QtGui.QImage(prediction['filename'])
            item = QtWidgets.QGraphicsPixmapItem(QtGui.QPixmap.fromImage(self._image))
            self._scene.addItem(item)
            self._scene.setSceneRect(item.boundingRect())
            self._rect = item.boundingRect()

            if mode in prediction.keys():
                self._overlay(prediction[mode])

        elif prediction['filetype'] == 'video':
            if 'cap' in prediction.keys():
                cap = prediction['cap']
            else:
                cap = Video(prediction['filename'])

            if 'frame' in prediction.keys():
                frame = prediction['frame']
            else:
                frame = 0

            video_frame = cap.get_frame(frame)
            if video_frame is not None:
                self._image = QtGui.QImage(video_frame,
                                           video_frame.shape[1], video_frame.shape[0],
                                           video_frame.shape[2] * video_frame.shape[1], QtGui.QImage.Format_RGB888)

                item = QtWidgets.QGraphicsPixmapItem(QtGui.QPixmap.fromImage(self._image))
                self._scene.addItem(item)
                self._scene.setSceneRect(item.boundingRect())
                self._rect = item.boundingRect()

                if 'prediction-by-frame' in prediction.keys():
                    self._overlay(prediction['prediction-by-frame'][frame])

        self.refit_display()

    def _overlay(self, prediction):
        if 'faces' in prediction.keys():
            self._overlay_faces(prediction['faces'])

    def _overlay_faces(self, faces):
        box_background = QtGui.QColor(255, 255, 255, 51)
        box_border = QtGui.QColor(0, 0, 0, 1)
        width = float(self._image.width())
        height = float(self._image.height())

        for face in faces:
            self._scene.addRect(face['x'] * width, face['y'] * height, face['w'] * width, face['h'] * height,
                                QtGui.QPen(QtGui.QBrush(box_border), 3.0),
                                QtGui.QBrush(box_background))

    def refit_display(self):
        if self._rect is None:
            self.resetTransform()
        else:
            self.fitInView(self._rect, QtCore.Qt.KeepAspectRatio)

    def resizeEvent(self, event: QtGui.QResizeEvent):
        self.refit_display()
        QtWidgets.QGraphicsView.resizeEvent(self, event)

    def reset(self):
        self._scene.clear()
        self._image = None
        self._current_video = None
        self._rect = None
        # TODO display context action