NegStation/negstation/widgets/framing_widget.py

import dearpygui.dearpygui as dpg
import numpy as np
import time
import scipy.ndimage as snd

from .stage_viewer_widget import PipelineStageViewer


class FramingWidget(PipelineStageViewer):
    name = "Framing"
    register = True
    has_pipeline_in = True
    has_pipeline_out = True

    def __init__(self, manager, logger):
        super().__init__(manager, logger)

        self.needs_publishing = False

        # Rotation line endpoints (canvas coords)
        self.rot_start = None  # (x, y)
        self.rot_end = None  # (x, y)
        self.angle = 0.0  # computed deskew angle

        # Crop rect endpoints (canvas coords)
        self.crop_start = None  # (x, y)
        self.crop_end = None    # (x, y)

        self.manager.bus.subscribe("img_clicked", self.on_click)
        self.manager.bus.subscribe("img_dragged", self.on_drag)
        self.manager.bus.subscribe("img_scrolled", self.on_scroll)
        self.manager.bus.subscribe("img_released", self.on_release)

    def create_pipeline_stage_content(self):
        super().create_pipeline_stage_content()

    def on_full_res_pipeline_data(self, img):
        if img is None:
            return
        img = self.rotate(img)
        self.publish_stage(self.crop(img))

    def on_pipeline_data(self, img):
        if img is None:
            return
        self.original_img = img.copy()
        self.needs_update = True

        # Apply transformation and publish
        self.img = self.rotate(img)
        self.publish_stage(self.crop(self.img))

    def update_texture(self, img):
        super().update_texture(img)

        # Draw rotation guide
        if self.rot_start and self.rot_end:
            p0 = self._pos_to_canvas(self.rot_start)
            p1 = self._pos_to_canvas(self.rot_end)
            dpg.draw_line(
                p1=p0,
                p2=p1,
                color=(255, 0, 0, 255),
                thickness=2,
                parent=self.drawlist,
            )

        # Draw crop rect
        if self.crop_start and self.crop_end:
            p0 = self._pos_to_canvas(self.crop_start)
            p1 = self._pos_to_canvas(self.crop_end)
            dpg.draw_rectangle(
                pmin=p0,
                pmax=p1,
                color=(255, 0, 255, 255),
                thickness=2,
                parent=self.drawlist,
            )

        if self.needs_publishing:
            img = self.crop(self.img)
            self.publish_stage(img)

    def _pos_to_canvas(self, img_pos):
        x, y = img_pos
        ix, iy = self.image_position
        iw, ih = self.img.shape[1], self.img.shape[0]
        sw, sh = self.scaled_size
        return (ix + x / iw * sw, iy + y / ih * sh)

    def on_click(self, data):
        if data.get("obj") is not self:
            return
        x, y = data.get("pos")
        if data["button"] == "left":
            self.rot_start = (x, y)
            self.rot_end = None
        elif data["button"] == "right":
            self.crop_start = (x,y)
            self.crop_end = None
        self.needs_update = True

    def on_drag(self, data):
        if data.get("obj") is not self:
            return
        x, y = data.get("pos")
        if data["button"] == "left":
            self.rot_end = (x, y)
        elif data["button"] == "right":
            self.crop_end = (x,y)
        self.needs_update = True

    def on_scroll(self, data):
        if data.get("obj") is not self:
            return

    def on_release(self, data):
        if data.get("obj") is not self:
            return
        if data["button"] == "left":
            # End of rotation line dragging
            dx = self.rot_end[0] - self.rot_start[0]
            dy = self.rot_end[1] - self.rot_start[1]
            self.angle += np.degrees(np.arctan2(dy, dx))
            # Do the rotation
            self.img = self.rotate(self.original_img)
            # Delete lines
            self.rot_start = self.rot_end = None
            self.needs_publishing = True
        elif data["button"] == "right":
            self.needs_publishing = True

        self.needs_update = True

    def rotate(self, img):
        img = snd.rotate(img, self.angle, reshape=False)
        return img

    def crop(self, img):
        """
        Crop `img` according to a rectangle drawn in the DISPLAY,
        by converting display‐coords to image‐coords via percentages.
        """
        # only crop if user has dragged out a box
        if not (self.crop_start and self.crop_end):
            return img

        # 1) get actual image dims
        h, w = img.shape[:2]

        # 2) get display info (set in update_texture)
        disp_h, disp_w, _ = self.img.shape

        # 3) unpack the two drag points (in screen coords)
        sx, sy = self.crop_start
        ex, ey = self.crop_end

        # 4) normalize into [0..1]
        nx0 = np.clip(sx / disp_w, 0.0, 1.0)
        ny0 = np.clip(sy / disp_h, 0.0, 1.0)
        nx1 = np.clip(ex / disp_w, 0.0, 1.0)
        ny1 = np.clip(ey / disp_h, 0.0, 1.0)

        # 5) map to image‐pixel coords
        x0 = int(nx0 * w)
        y0 = int(ny0 * h)
        x1 = int(nx1 * w)
        y1 = int(ny1 * h)

        # 6) sort & clamp again just in case
        x0, x1 = sorted((max(0, x0), min(w, x1)))
        y0, y1 = sorted((max(0, y0), min(h, y1)))

        # 7) avoid zero‐area
        if x0 == x1 or y0 == y1:
            return img

        return img[y0:y1, x0:x1]