SwissMX/CustomROI.py

import logging

from PyQt5 import QtGui, QtCore
from PyQt5.QtGui import QColor, QTransform
from PyQt5.QtWidgets import QMenu, QAction, QSpinBox, QMenu

#from GenericDialog import GenericDialog #ZAC: orig. code
from pyqtgraph import ROI, Point, RectROI, mkPen
from pyqtgraph import functions as fn

from app_config import settings

from math import *

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


class Retangulo(ROI):
    """
    Rectangular ROI subclass with a single scale handle at the top-right corner.

    ============== =============================================================
    **Arguments**
    pos            (length-2 sequence) The position of the ROI origin.
                   See ROI().
    size           (length-2 sequence) The size of the ROI. See ROI().
    centered       (bool) If True, scale handles affect the ROI relative to its
                   center, rather than its origin.
    sideScalers    (bool) If True, extra scale handles are added at the top and
                   right edges.
    \**args        All extra keyword arguments are passed to ROI()
    ============== =============================================================

    """

    def __init__(self, pos, size, centered=False, sideScalers=False, **args):
        ROI.__init__(self, pos, size, **args)
        if centered:
            center = [0.5, 0.5]
        else:
            center = [0, 0]

        self.addScaleRotateHandle([0, 0.5], [1, 0.5])
        self.addScaleHandle([1, 1], center)
        if sideScalers:
            self.addScaleHandle([1, 0.5], [center[0], 0.5])
            self.addScaleHandle([0.5, 1], [0.5, center[1]])


class Grid(ROI):
    """a grid"""

    def __init__(
        self,
        x_step,
        y_step,
        x_offset,
        y_offset,
        grid_index=0,
        parent=None,
        gonio_xy=(0, 0),
        **kargs
    ):
        self._grid_index = grid_index
        self._size = None
        self._pos = None
        self._shape = None
        ROI.__init__(self, (500, 500), (300, 300), **kargs)
        self.parent = parent

        self.addScaleHandle([1, 1], [0, 0])
        self.addScaleHandle([0, 0], [1, 1])
        self.addScaleRotateHandle([1, 0], [0, 0])
        # self.addScaleRotateHandle([1, 0.5], [0.5, 0.5])

        self._steps = x_step / 1000., y_step / 1000.
        self._offsets = x_offset / 1000, y_offset / 1000
        self._beam_position = parent.get_beam_mark_on_camera_xy()
        self.set_gonio_xy(*gonio_xy)

        self.setToolTip("grid")

        parent.pixelsPerMillimeter.connect(self.update_ppm)
        parent.beamCameraCoordinatesChanged.connect(self.update_beam_position)
        self.enable_stage_tracking()

        self._ppm = parent.getPpm()
        self._mm_size = 300 / self._ppm, 300 / self._ppm
        self._original_ppm = self._ppm  # marker only works at original zoom
        # because zoom axis is not aligned
        # with viewing axis
        self._fx, self._fy = self.parent.getFastX(), self.parent.getFastY()
        self.invertible = True
        self.sigRegionChanged.connect(self.invalidate)
        self.sigRegionChangeFinished.connect(self.calculate_gonio_xy)
        self.scaleSnap = False
        self.aspectLocked = False
        self.translatable = True
        self.rotateAllowed = False
        self.removable = True

        self.follow_stage()
        # self.calculate_gonio_xy()

    def enable_stage_tracking(self):
        self.parent.fast_x_position.connect(self.follow_x)
        self.parent.fast_y_position.connect(self.follow_y)

    def disable_stage_tracking(self):
        self.parent.fast_x_position.disconnect(self.follow_x)
        self.parent.fast_y_position.disconnect(self.follow_y)

    def follow_x(self, v):
        self.follow_stage(fx=v)

    def follow_y(self, v):
        self.follow_stage(fy=v)

    def removeClicked(self):
        self.parent.pixelsPerMillimeter.disconnect(self.update_ppm)
        self.parent.beamCameraCoordinatesChanged.disconnect(self.update_beam_position)
        self.parent.fast_x_position.disconnect(self.follow_x)
        self.parent.fast_y_position.disconnect(self.follow_y)
        super(Grid, self).removeClicked()

    def calculate_gonio_xy(self):
        self.update_grid_pos()
        state = self.getState()
        angle = state["angle"]
        w, h = state["size"]
        w, h = 1000 * w / self._ppm, 1000 * h / self._ppm
        self._mm_size = w / 1000, h / 1000
        self.calculate_targets()

    def update_ppm(self, ppm):
        self._ppm = ppm
        self.update_grid_size(ppm)
        self.update_grid_pos()
        self.stateChanged()

    def update_grid_size(self, ppm, finish=False):
        mmw, mmh = self._mm_size
        w, h = mmw * ppm, mmh * ppm
        self.setSize((w, h), update=False)

    def follow_stage(self, fx=None, fy=None):
        self._fx, self._fy = self.parent.getFastX(), self.parent.getFastY()
        self.update_grid_pos()
        self.stateChanged()

    def set_gonio_xy(self, x, y):
        self._gx, self._gy = x, y

    def update_grid_pos(self):
        idx = self._grid_index
        ppm = self._ppm
        bx, by = self.parent.get_beam_mark_on_camera_xy()
        gx, gy = self._gx, self._gy  # marker position in gonio coords
        fx, fy = self._fx, self._fy
        dx, dy = (
            ppm * (fx - gx),
            ppm * (fy - gy),
        )  # distance to beam in pixels for new ppm
        x, y = bx - dx, by + dy
        self.setPos(x, y, update=False)
        mmw, mmh = self._mm_size

    def calculate_targets(self):
        state = self.getState()
        angle = radians(state["angle"])
        cosa = cos(angle)
        sina = sin(angle)
        w, h = state["size"]
        xs, ys = self._steps
        w, h = w / self._ppm, h / self._ppm  # mm
        gx, gy = self._gx, self._gy
        nx = int(abs(w) / xs)
        ny = int(abs(h) / ys)
        points = []
        dx = xs
        dy = ys
        for yi in range(ny):
            ly = yi * dy
            y = gy + ly
            for xi in range(nx):
                lx = xi * dx
                x = gx + lx
                l = sqrt((lx) ** 2 + (ly) ** 2)
                try:
                    alpha = acos((lx) / l)
                except:
                    alpha = 0
                x2 = gx + l * cos(alpha + angle)
                y2 = gy - l * sin(alpha + angle)
                # print(
                #     "l={:8.3f} alpha={:.3f} xy={:8.3f}{:8.3f} xy2={:8.3f}{:8.3f}".format(
                #         l, degrees(alpha), x, y, x2, y2
                #     )
                # )

                points.append((x2, y2))
        self._grid_dimensions = (nx, ny)
        self._points = points
        self.parent.gridUpdated.emit(self._grid_index)

    def get_sorted_grid_targets(self):
        pass

    def get_grid_targets(self):
        return self._points

    def update_beam_position(self, pos):
        self._beam_position = pos

    def invalidate(self):
        state = self.getState()
        x, y = state["pos"]
        w, h = state["size"]
        self._mm_size = w / self._ppm, h / self._ppm
        fx, fy = self.camera2gonio(x, y)
        self.set_gonio_xy(fx, fy)
        self._shape = None
        self.prepareGeometryChange()

    def boundingRect(self):
        return self.shape().boundingRect()

    def shape(self):
        w, h = self.getState()["size"]
        if self._shape is None:
            radius = self.getState()["size"][1]
            p = QtGui.QPainterPath()
            p.addRect(0, 0, w, h)
            stroker = QtGui.QPainterPathStroker()
            stroker.setWidth(10)
            outline = stroker.createStroke(p)
            self._shape = outline
        return self._shape

    def paint(self, p, *args):
        state = self.getState()
        w, h = state["size"]
        p.setRenderHint(QtGui.QPainter.Antialiasing)
        p.setPen(fn.mkPen(width=5, color=[200, 200, 100]))
        p.drawRect(0, 0, w, h)

        p.setPen(fn.mkPen(width=3, color=[200, 100, 100]))
        ppm = self._ppm
        xs, ys = self._steps
        w1, h1 = w / ppm, h / ppm  # mm
        nx = int(abs(w1) / xs)
        ny = int(abs(h1) / ys)
        sw = (xs) * ppm
        sh = (ys) * ppm
        # print((nx,ny), (w1, h1), (sw, sh))
        r = 10
        a, b = state["pos"]
        if nx * ny > 1000:
            return
        for y in range(ny):
            for x in range(nx):
                p.drawEllipse((x * sw) - r, (y * sh) - r, 2 * r, 2 * r)

    def camera2gonio(self, x, y):
        bx, by = self.parent.get_beam_mark_on_camera_xy()
        ppm = self._ppm
        dx, dy = (x - bx) / ppm, (y - by) / ppm
        fx, fy = self._fx + dx, self._fy - dy
        return fx, fy

    def gonio2camera(self, x, y):
        ppm = self._ppm
        bx, by = self.parent.get_beam_mark_on_camera_xy()
        gx, gy = self._gx, self._gy  # marker position in gonio coords
        fx, fy = self._fx, self._fy
        dx, dy = (
            ppm * (fx - gx),
            ppm * (fy - gy),
        )  # distance to beam in pixels for new ppm
        x, y = bx - dx, by + dy
        return x, y


class BeamMark(ROI):
    """A crosshair ROI whose position is at the center of the crosshairs. By default, it is scalable, rotatable and translatable."""

    def __init__(self, pos=None, size=None, parent=None, **kargs):
        if size is None:
            size = [20, 20]
        if pos is None:
            pos = [0, 0]
        self._size = size
        self._pos = pos
        self._shape = None
        ROI.__init__(self, pos, size, **kargs)

        parent.pixelsPerMillimeter.connect(self.my_update_ppm)
        parent.beamCameraCoordinatesChanged.connect(lambda x, y: self.setPos((x, y)))

        self._scale_handler = None
        self.sigRegionChanged.connect(self.invalidate)
        self.aspectLocked = False
        self.translatable = False
        self.rotateAllowed = False
        self.invertible = False
        self.removable = False
        self.scaleSnap = True
        self.snapSize = 1
        self.sigRegionChangeFinished.connect(self.show_size)
        self.flip_ud()

    def show_size(self):
        try:
            ppm = self._ppm
        except:
            return
        w, h = self.getState()["size"]
        # convert to micrometer
        nw, nh = (w * 1000) / ppm, (h * 1000) / ppm
        logger.debug("persisting new beam marker size: {}".format((nw, nh)))
        settings.setValue("beam/size", (nw, nh))

    def toggle_handle(self):
        if self._scale_handler is None:
            logger.debug("adding beam marker resize handle")
            self._scale_handler = self.addScaleHandle([0.5, 0.5], [0, 0])
        else:
            logger.debug("removing beam marker resize handle")
            self.removeHandle(self._scale_handler)
            self._scale_handler = None

    def get_camera_beam_position(self):
        return self.getState()["pos"]

    def set_beam_size_marker_dialog(self):
        w, h = settings.value("beam/size")
        d = GenericDialog(
            title="Beamsize",
            message="Enter the size of the beam in microns",
            inputs={
                "width": ("Width (\u00B5)", int(w), QSpinBox()),
                "height": ("Height (\u00B5)", int(h), QSpinBox()),
            },
        )
        if d.exec():
            results = d.results
            logger.info("Updating beamsize to {}".format(results))
            w, h = results["width"], results["height"]
            logger.debug("types {}".format(type(w)))
            settings.setValue("beam/size", (w, h))
            self.set_beam_size((w, h))

    def flip_ud(self):
        """flip up-down"""
        t = self.transform()
        m11 = t.m11()
        m12 = t.m12()
        m13 = t.m13()
        m21 = t.m21()
        m22 = t.m22()
        m23 = t.m23()
        m31 = t.m31()
        m32 = t.m32()
        m33 = t.m33()
        t.setMatrix(m11, m12, m13, m21, -m22, m23, m31, m32, m33)
        self.setTransform(t)

    def set_beam_size(self, size=None):
        if size is not None:
            self._size = size
        ppm = self._ppm
        w, h = self._size
        nw, nh = (w / 1000) * ppm, (h / 1000) * ppm
        self.setSize((nw, nh))
        self.setToolTip("Beamsize (W x H): {:.1f}\u00B5 x {:.1f}\u00B5".format(w, h))

    def my_update_ppm(self, ppm):
        self._ppm = ppm
        self.set_beam_size()

    def invalidate(self):
        self._shape = None
        self.prepareGeometryChange()

    def boundingRect(self):
        return self.shape().boundingRect()

    def shape(self):
        w, h = self.size()
        if self._shape is None:
            p = QtGui.QPainterPath()
            p.moveTo(Point(0, -h))
            p.lineTo(Point(0, h))
            p.moveTo(Point(-w, 0))
            p.lineTo(Point(w, 0))
            #
            p.moveTo(Point(-w, -h))
            p.lineTo(Point(w, -h))
            p.lineTo(Point(w, h))
            p.lineTo(Point(-w, -h))
            stroker = QtGui.QPainterPathStroker()
            outline = stroker.createStroke(p)
            self._shape = outline

        return self._shape

    def paint(self, p, *args):
        w, h = self.getState()["size"]

        v1, v2 = -(h * 0.8) / 2, (h * 0.8) / 2
        h1, h2 = -(w * 0.8) / 2, (w * 0.8) / 2

        p.setRenderHint(QtGui.QPainter.Antialiasing)
        p.setPen(fn.mkPen(width=3, color=[200, 100, 100]))
        p.drawLine(Point(0, v1), Point(0, v2))
        p.drawLine(Point(h1, 0), Point(h2, 0))

        p.setPen(fn.mkPen(width=3, color=[200, 200, 100]))
        p.drawRect(-w / 2, -h / 2, w, h)
        # p.drawText(-w, -h, '{:.0f}x{:.0f}'.format(*self._size))


class CrystalCircle(ROI):
    """A crosshair ROI whose position is at the center of the crosshairs. By default, it is scalable, rotatable and translatable."""

    def __init__(
        self,
        pos=None,
        size=None,
        parent=None,
        gonio_xy=(0, 0),
        model_row_index=1,
        is_fiducial=False,
        ppm=False,
        **kargs
    ):
        if size is None:
            size = [20, 20]
        if pos is None:
            pos = [0, 0]
        self._size = size
        self._shape = None
        self.parent = parent
        self._model_row_index = model_row_index
        self._is_fiducial = is_fiducial
        ROI.__init__(self, pos, size, **kargs)

        self._beam_position = parent.get_beam_mark_on_camera_xy()
        self.set_gonio_xy(*gonio_xy)

        self.setToolTip(f"{'fiducial' if is_fiducial else 'crystal'} {model_row_index}")

        # parent.pixelsPerMillimeter.connect(self.update_pixels_per_millimeter)
        # parent.beamCameraCoordinatesChanged.connect(self.update_beam_position)
        # parent.fast_x_position.connect(self.follow_x)
        # parent.fast_y_position.connect(self.follow_y)
        self._ppm = ppm
        self._original_ppm = ppm  # marker only works at original zoom
        # because zoom axis is not aligned
        # with viewing axis

        # self.calculate_gonio_xy()
        self.sigRegionChanged.connect(self.invalidate)
        # self.sigRegionChangeFinished.connect(self.calculate_gonio_xy)
        # self.sigHoverEvent.connect(self.my_hovering)
        self.aspectLocked = False
        self.translatable = True
        self.rotateAllowed = False
        self.invertible = False
        self.deletable = False

    def enable_stage_tracking(self):
        self.parent.fast_x_position.connect(self.follow_x)
        self.parent.fast_y_position.connect(self.follow_y)

    def disable_stage_tracking(self):
        self.parent.fast_x_position.disconnect(self.follow_x)
        self.parent.fast_y_position.disconnect(self.follow_y)

    def follow_x(self, v):
        self.follow_stage(fx=v)

    def follow_y(self, v):
        self.follow_stage(fy=v)

    def my_hovering(self, event):
        print("hovering: {}".format(self.mouseHovering))

    def get_model_row(self):
        return self._model_row_index

    def update_ppm(self, ppm):
        self._ppm = ppm
        self.recalc_camera_position()

    def follow_stage(self, fx=None, fy=None):
        self._fx, self._fy = self.parent.getFastX(), self.parent.getFastY()
        self.recalc_camera_position()

    def set_gonio_xy(self, x, y):
        self._gx, self._gy = x, y

    def disconnect_signals(self):
        self.parent.pixelsPerMillimeter.disconnect(self.update_ppm)
        self.parent.beamCameraCoordinatesChanged.disconnect(self.update_beam_position)
        self.parent.fast_x_position.disconnect(self.follow_x)
        self.parent.fast_y_position.disconnect(self.follow_y)

    def reconnect_signals(self):
        self.parent.pixelsPerMillimeter.connect(self.update_ppm)
        self.parent.beamCameraCoordinatesChanged.connect(self.update_beam_position)
        self.parent.fast_x_position.connect(self.follow_x)
        self.parent.fast_y_position.connect(self.follow_y)

    def removeClicked(self):
        self.disconnect_signals()
        super(CrystalCircle, self).removeClicked()

    def recalc_camera_position(self):
        idx = self._model_row_index
        ppm = self._ppm
        bx, by = self.parent.get_beam_mark_on_camera_xy()
        gx, gy = self._gx, self._gy  # marker position in gonio coords
        fx, fy = self._fx, self._fy
        # distance to beam in pixels for new ppm
        dx, dy = (ppm * (fx - gx), ppm * (fy - gy))
        x, y = bx - dx, by + dy
        # logger.debug(f"recalc {idx}: cam = {(x,y)}, gonio = {(gx, gy)}")
        self.setPos(x, y)

    def update_beam_position(self, pos):
        self._beam_position = pos

    def show_size(self):
        try:
            ppm = self._ppm
        except:
            return
        w, h = self.getState()["size"]
        # convert to micrometer
        nw, nh = (w * 1000) / ppm, (h * 1000) / ppm

    def invalidate(self):
        self._shape = None
        self.prepareGeometryChange()

    def boundingRect(self):
        return self.shape().boundingRect()

    def shape(self):
        if self._shape is None:
            radius = self.getState()["size"][1]
            p = QtGui.QPainterPath()
            p.moveTo(Point(0, -radius))
            p.lineTo(Point(0, radius))
            p.moveTo(Point(-radius, 0))
            p.lineTo(Point(radius, 0))
            stroker = QtGui.QPainterPathStroker()
            stroker.setWidth(10)
            outline = stroker.createStroke(p)
            self._shape = outline

        return self._shape

    def paint(self, p, *args):
        if abs(self._ppm - self._original_ppm) < 1.:
            if self._is_fiducial:
                pen = mkPen(color="r", width=2)
            else:
                pen = mkPen(color="y", width=2)
        else:
            pen = mkPen(color="#ffffaa88", width=2)
        x, y = self.getState()["size"]
        p.setRenderHint(QtGui.QPainter.Antialiasing)
        p.setPen(pen)

        p.drawLine(Point(0, -x), Point(0, x))
        p.drawLine(Point(-x, 0), Point(x, 0))
        p.drawEllipse(-x, -y, 2 * x, 2 * y)


class Crosshair(ROI):
    """A crosshair ROI whose position is at the center of the crosshairs. By default, it is scalable, rotatable and translatable."""

    def __init__(self, pos=None, size=None, **kargs):
        if size is None:
            size = [10, 10]
        if pos is None:
            pos = [0, 0]
        self._shape = None
        ROI.__init__(self, pos, size, **kargs)

        self.sigRegionChanged.connect(self.invalidate)
        self.aspectLocked = True

    def invalidate(self):
        self._shape = None
        self.prepareGeometryChange()

    def boundingRect(self):
        return self.shape().boundingRect()

    def shape(self):
        if self._shape is None:
            radius = self.getState()["size"][1]
            p = QtGui.QPainterPath()
            p.moveTo(Point(0, -radius))
            p.lineTo(Point(0, radius))
            p.moveTo(Point(-radius, 0))
            p.lineTo(Point(radius, 0))
            # p = self.mapToDevice(p)
            stroker = QtGui.QPainterPathStroker()
            stroker.setWidth(10)
            outline = stroker.createStroke(p)
            # self._shape = self.mapFromDevice(outline)
            self._shape = outline

        return self._shape

    def paint(self, p, *args):
        radius = self.getState()["size"][1]
        p.setRenderHint(QtGui.QPainter.Antialiasing)
        p.setPen(self.currentPen)

        p.drawLine(Point(0, -radius), Point(0, radius))
        p.drawLine(Point(-radius, 0), Point(radius, 0))