Annotation WIP

annotation.py

"""
A very basic example of a image annotation tool using a QWidget as a viewport.

This implements
    - Viewport panning
    - Viewport zooming (maintaing Point Of Interest)
    - Basic draggable handles / gizmos
"""

from PySide6 import QtCore, QtGui, QtWidgets


class DraggableMixin:
    def __init__(self):
        self._pos = QtCore.QPoint(0, 0)

    def translate(self, point: QtCore.QPoint):
        self._pos += point

    def rect(self) -> QtCore.QRect:
        raise NotImplementedError()


class Handle(DraggableMixin):
    def __init__(self):
        super().__init__()

    def rect(self) -> QtCore.QRect:
        return QtCore.QRect(self._pos, QtCore.QSize(30, 30))

    def paint(self, painter: QtGui.QPainter):
        painter.setPen(QtCore.Qt.red)
        painter.setBrush(QtCore.Qt.NoBrush)
        painter.drawRect(self.rect())


class Annotation:
    def __init__(self):
        self.handles = []

    def rect(self) -> QtCore.QRect:
        raise NotImplementedError()

    def draw(self, painter: QtGui.QPainter):
        painter.setPen(QtCore.Qt.red)
        painter.setBrush(QtCore.Qt.NoBrush)
        painter.drawRect(self.rect())


class RectAnnotation(Annotation):
    def __init__(self):
        super().__init__()

        self.handles.append(Handle())
        self.handles.append(Handle())

    def rect(self) -> QtCore.QRect:
        return QtCore.QRect(self.handles[0].rect().topLeft(), self.handles[1].rect().bottomRight())

    def draw(self, painter: QtGui.QPainter):
        painter.setPen(QtCore.Qt.red)
        painter.setBrush(QtCore.Qt.NoBrush)
        painter.drawRect(self.rect())


class Viewport(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMouseTracking(True)

        self._transform = QtGui.QTransform()

        self.bg_brush = QtWidgets.QApplication.palette().brush(
            QtGui.QPalette.ColorGroup.Active,
            QtGui.QPalette.ColorRole.Window,

        )

        self._last_mouse_pos = QtCore.QPoint(0, 0)

        self._bg_pixmap: QtGui.QPixmap = None
        self._bg_pixmap_rect: QtCore.QRect = None

        self.h1 = Handle()
        self.h2 = Handle()

        self.h1.translate(QtCore.QPoint(-100, -100))
        self.h2.translate(QtCore.QPoint(100, 100))

        self._annotation = RectAnnotation()

        self._handles = self._annotation.handles

    def setPixmap(self, pixmap: QtGui.QPixmap):
        self._bg_pixmap = pixmap
        self._bg_pixmap_rect = pixmap.rect()
        self._bg_pixmap_rect.moveCenter(QtCore.QPoint(0, 0))
        self.update()

    def getViewportTransform(self) -> QtGui.QTransform:
        transform = QtGui.QTransform()
        transform.translate(self.width() / 2, self.height() / 2)
        transform *= self._transform
        return transform

    def paintEvent(self, event: QtGui.QPaintEvent):
        painter = QtGui.QPainter(self)
        painter.setRenderHint(QtGui.QPainter.Antialiasing)

        painter.setBrush(self.bg_brush)
        painter.drawRect(event.rect())

        view_transform = QtGui.QTransform()
        view_transform.translate(self.width() / 2, self.height() / 2)
        view_transform *= self._transform

        painter.setTransform(view_transform)

        if self._bg_pixmap:
            painter.drawPixmap(self._bg_pixmap_rect, self._bg_pixmap)

        self._annotation.draw(painter)

        if self._handles:
            for handle in self._handles:
                handle.paint(painter)

        mouse_pos = self.mapFromGlobal(QtGui.QCursor.pos())
        mouse_view_pos = view_transform.inverted()[0].map(mouse_pos)

        painter.setPen(QtGui.QPen(QtCore.Qt.red))
        painter.drawEllipse(mouse_view_pos, 5, 5)

        painter.end()

    def pan(self, x, y, update=True):
        self._transform.translate(x, y)
        if update:
            self.update()

    def zoom(self, factor, update=True):
        self._transform.scale(factor, factor)
        if update:
            self.update()

    def mousePressEvent(self, event: QtGui.QMouseEvent):
        self._last_mouse_pos = event.position()
        if event.button() == QtCore.Qt.MouseButton.MiddleButton:
            self.setCursor(QtCore.Qt.CursorShape.ClosedHandCursor)

    def mouseMoveEvent(self, event: QtGui.QMouseEvent):
        mouse_world_pos = self.viewToScene(event.pos())

        if event.buttons() == QtCore.Qt.MouseButton.MiddleButton:
            delta = event.position() - self._last_mouse_pos

            delta /= self._transform.m11()

            self._last_mouse_pos = event.position()
            self.pan(delta.x(), delta.y(), update=True)

        elif event.buttons() == QtCore.Qt.MouseButton.NoButton:
            for handle in self._handles:
                if handle.rect().contains(mouse_world_pos):
                    self.setCursor(QtCore.Qt.CursorShape.OpenHandCursor)
                    break
            else:
                self.setCursor(QtCore.Qt.CursorShape.ArrowCursor)

        elif event.buttons() == QtCore.Qt.MouseButton.LeftButton:
            for handle in self._handles:
                if handle.rect().contains(mouse_world_pos):
                    handle.translate(mouse_world_pos - handle.rect().center())
                    self.update()
                    break

        self.update()

    def mouseReleaseEvent(self, event: QtGui.QMouseEvent):
        self.setCursor(QtCore.Qt.CursorShape.ArrowCursor)

    def viewToScene(self, point: QtCore.QPoint) -> QtCore.QPoint:
        return self.getViewportTransform().inverted()[0].map(point)

    def wheelEvent(self, event: QtGui.QWheelEvent):
        initial_pos = self.viewToScene(event.position().toPoint())

        delta = event.angleDelta().y()

        if delta > 0 and self._transform.m11() < 5:
            self.zoom(1.05, update=False)
        elif delta < 0 and self._transform.m11() > 0.5:
            self.zoom(0.95, update=False)

        final_pos = self.viewToScene(event.position().toPoint())
        delta = final_pos - initial_pos
        self.pan(delta.x(), delta.y(), update=True)


if __name__ == '__main__':
    import os
    import urllib.request
    import tempfile

    _, filename = tempfile.mkstemp(suffix=".jpg")
    urllib.request.urlretrieve("https://picsum.photos/200/300", filename)
    try:
        app = QtWidgets.QApplication([])
        pixmap = QtGui.QPixmap(filename)
        view = Viewport()
        view.setPixmap(pixmap)
        view.show()
        app.exec()
    finally:
        os.remove(filename)