ScanSlicer/scanslicer/canvas_state.py

import numpy as np
from PySide6.QtCore import QPointF, QRectF

from .slices import PhotoSlice, polys_iou


# CanvasStateMixin: class definition for ScanSlicer behavior.
class CanvasStateMixin:
    # fit_image_to_view: function definition for ScanSlicer behavior.
    def fit_image_to_view(self):
        if self.image is None:
            return

        view_w = max(1, float(self.width()))
        view_h = max(1, float(self.height()))
        image_w = max(1, float(self.image.width()))
        image_h = max(1, float(self.image.height()))

        self.zoom = min(view_w / image_w, view_h / image_h) * self.AUTO_FIT_SCALE

        shown_w = image_w * self.zoom
        shown_h = image_h * self.zoom
        self.pan_offset = QPointF((view_w - shown_w) / 2.0, (view_h - shown_h) / 2.0)

    # reset_view: function definition for ScanSlicer behavior.
    def reset_view(self):
        if self.image is None:
            return False

        self.user_transformed_view = False
        self.fit_image_to_view()
        self.update()
        return True

    # set_image: function definition for ScanSlicer behavior.
    def set_image(self, image, new_image=False):
        if self.image is None or new_image:
            self.user_transformed_view = False
            self.slices = []
            self.set_active_slice(None)
            self.box_dragging = None
        self.image = image
        if self.image is not None and not self.user_transformed_view:
            self.fit_image_to_view()
        self.update()

    # _clamp_slice_to_image: function definition for ScanSlicer behavior.
    def _clamp_slice_to_image(self, slice_obj):
        if slice_obj is None:
            return None
        slice_obj.bbox = self._clamp_bbox_to_image(slice_obj.bbox)
        return slice_obj

    # _coerce_slice_in_bounds: function definition for ScanSlicer behavior.
    def _coerce_slice_in_bounds(self, bbx):
        if isinstance(bbx, PhotoSlice):
            return self._clamp_slice_to_image(bbx)
        return PhotoSlice(self._clamp_bbox_to_image(bbx))

    # add_bbox: function definition for ScanSlicer behavior.
    def add_bbox(self, bbx):
        self.slices.append(self._coerce_slice_in_bounds(bbx))
        self.set_active_slice(len(self.slices) - 1)
        self.update()

    @staticmethod
    # _aabb_overlaps: function definition for ScanSlicer behavior.
    def _aabb_overlaps(poly_a, poly_b):
        a = np.array(poly_a, dtype=np.float32).reshape(-1, 2)
        b = np.array(poly_b, dtype=np.float32).reshape(-1, 2)
        a_min_x = float(np.min(a[:, 0]))
        a_max_x = float(np.max(a[:, 0]))
        a_min_y = float(np.min(a[:, 1]))
        a_max_y = float(np.max(a[:, 1]))
        b_min_x = float(np.min(b[:, 0]))
        b_max_x = float(np.max(b[:, 0]))
        b_min_y = float(np.min(b[:, 1]))
        b_max_y = float(np.max(b[:, 1]))
        return not (
            a_max_x < b_min_x
            or b_max_x < a_min_x
            or a_max_y < b_min_y
            or b_max_y < a_min_y
        )

    # update_bboxes: function definition for ScanSlicer behavior.
    def update_bboxes(self, bbxs=None):
        if bbxs is not None:
            active_slice = None
            if self.active_slice_index is not None and 0 <= self.active_slice_index < len(self.slices):
                active_slice = self.slices[self.active_slice_index]

            merged_slices = []
            for sl in self.slices:
                if sl.locked:
                    merged_slices.append(self._clamp_slice_to_image(sl))
            for box in bbxs:
                box = self._clamp_bbox_to_image(box)
                overlaps_locked = False
                for sl in merged_slices:
                    if not self._aabb_overlaps(box, sl.bbox):
                        continue
                    if polys_iou(box, sl.bbox) > 0.3:
                        overlaps_locked = True
                        break
                if not overlaps_locked:
                    merged_slices.append(PhotoSlice(box))

            self.slices = merged_slices
            if active_slice is not None and active_slice in self.slices:
                self.set_active_slice(self.slices.index(active_slice))
            else:
                self.set_active_slice(None)

        self.update()

    # set_active_slice: function definition for ScanSlicer behavior.
    def set_active_slice(self, index):
        if index is None:
            normalized = None
        elif 0 <= int(index) < len(self.slices):
            normalized = int(index)
        else:
            normalized = None

        if normalized == self.active_slice_index:
            return

        self.active_slice_index = normalized
        self.active_slice_changed.emit(-1 if normalized is None else normalized)
        self.update()

    # get_active_slice: function definition for ScanSlicer behavior.
    def get_active_slice(self):
        if self.active_slice_index is None:
            return None
        if not (0 <= self.active_slice_index < len(self.slices)):
            self.set_active_slice(None)
            return None
        return self.slices[self.active_slice_index]

    # set_active_slice_locked: function definition for ScanSlicer behavior.
    def set_active_slice_locked(self, locked):
        active_slice = self.get_active_slice()
        if active_slice is None:
            return False
        active_slice.toggle_locked(bool(locked))
        self.update()
        return True

    # set_all_slices_locked: function definition for ScanSlicer behavior.
    def set_all_slices_locked(self, locked):
        for s in self.slices:
            s.toggle_locked(bool(locked))
        self.update()
        return len(self.slices)

    # are_all_slices_locked: function definition for ScanSlicer behavior.
    def are_all_slices_locked(self):
        if len(self.slices) == 0:
            return False
        return all(s.locked for s in self.slices)

    # remove_active_slice: function definition for ScanSlicer behavior.
    def remove_active_slice(self):
        if self.active_slice_index is None:
            return False
        if not (0 <= self.active_slice_index < len(self.slices)):
            self.set_active_slice(None)
            return False

        remove_index = self.active_slice_index
        self.slices.pop(remove_index)

        if len(self.slices) == 0:
            self.set_active_slice(None)
        else:
            new_index = min(remove_index, len(self.slices) - 1)
            # Force refresh even if the numeric index is unchanged.
            self.active_slice_index = None
            self.set_active_slice(new_index)

        self.update()
        return True

    # squarize_active_slice: function definition for ScanSlicer behavior.
    def squarize_active_slice(self):
        active_slice = self.get_active_slice()
        if active_slice is None:
            return False
        active_slice.squarize()
        self.update()
        return True

    # reset_active_slice_rotation: function definition for ScanSlicer behavior.
    def reset_active_slice_rotation(self):
        active_slice = self.get_active_slice()
        if active_slice is None:
            return False
        active_slice.reset_rotation()
        active_slice.bbox = self._clamp_bbox_to_image(active_slice.bbox)
        self.update()
        return True

    # set_rotate_handles_mode: function definition for ScanSlicer behavior.
    def set_rotate_handles_mode(self, enabled):
        self.rotate_handles_mode = bool(enabled)

    # _image_rect_view: function definition for ScanSlicer behavior.
    def _image_rect_view(self):
        if self.image is None:
            return QRectF()
        return QRectF(
            self.pan_offset.x(),
            self.pan_offset.y(),
            self.image.width() * self.zoom,
            self.image.height() * self.zoom,
        )

    # _view_to_image: function definition for ScanSlicer behavior.
    def _view_to_image(self, p):
        return np.array(
            [
                (p.x() - self.pan_offset.x()) / self.zoom,
                (p.y() - self.pan_offset.y()) / self.zoom,
            ],
            dtype=np.float64,
        )

    # _clamp_view_pos_to_canvas: function definition for ScanSlicer behavior.
    def _clamp_view_pos_to_canvas(self, p):
        max_x = max(0.0, float(self.width() - 1))
        max_y = max(0.0, float(self.height() - 1))
        return QPointF(
            min(max(p.x(), 0.0), max_x),
            min(max(p.y(), 0.0), max_y),
        )

    # _clamp_image_point: function definition for ScanSlicer behavior.
    def _clamp_image_point(self, point):
        if self.image is None:
            return np.array(point, dtype=np.float64)

        max_x = max(0.0, float(self.image.width() - 1))
        max_y = max(0.0, float(self.image.height() - 1))
        p = np.array(point, dtype=np.float64).reshape(2)
        p[0] = min(max(p[0], 0.0), max_x)
        p[1] = min(max(p[1], 0.0), max_y)
        return p

    # _clamp_bbox_to_image: function definition for ScanSlicer behavior.
    def _clamp_bbox_to_image(self, bbox):
        if self.image is None:
            return np.array(bbox, dtype=np.float64)

        max_x = max(0.0, float(self.image.width() - 1))
        max_y = max(0.0, float(self.image.height() - 1))
        pts = np.array(bbox, dtype=np.float64).reshape(4, 2)
        pts[:, 0] = np.clip(pts[:, 0], 0.0, max_x)
        pts[:, 1] = np.clip(pts[:, 1], 0.0, max_y)
        return pts

    # _image_to_view: function definition for ScanSlicer behavior.
    def _image_to_view(self, p):
        return QPointF(
            p[0] * self.zoom + self.pan_offset.x(),
            p[1] * self.zoom + self.pan_offset.y(),
        )