imageGraphicsView.py

from scipy.signal import savgol_filter
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsPixmapItem
from PyQt5.QtGui import QPixmap, QColor
from PyQt5.QtCore import Qt, QRectF, QPoint
import math
import numpy as np
from panZoomGraphicsView import PanZoomGraphicsView
from labeledPointItem import LabeledPointItem
from find_path import find_path


class ImageGraphicsView(PanZoomGraphicsView):
    """
    A custom QGraphicsView for displaying and interacting with an image.

    This class extends PanZoomGraphicsView to provide additional functionality
    for loading images, adding labeled anchor points, and computing paths
    between points based on a cost image.
    """

    def __init__(self, parent=None):
        super().__init__(parent)
        self.scene = QGraphicsScene(self)
        self.setScene(self.scene)

        # Image display
        self.image_item = QGraphicsPixmapItem()
        self.scene.addItem(self.image_item)

        self.anchor_points = []    
        self.point_items = []      
        self.full_path_points = [] 
        self._full_path_xy = []    

        self.dot_radius = 4
        self.path_radius = 1
        self.radius_cost_image = 2
        self._img_w = 0
        self._img_h = 0

        self._mouse_pressed = False
        self._press_view_pos = None
        self._drag_threshold = 5
        self._was_dragging = False
        self._dragging_idx = None
        self._drag_offset = (0, 0)
        self._drag_counter = 0

        # Cost images
        self.cost_image_original = None
        self.cost_image = None

        # Rainbow toggle => start with OFF
        self._rainbow_enabled = False

        # Smoothing parameters
        self._savgol_window_length = 7

    def set_rainbow_enabled(self, enabled: bool):
        """Enable rainbow coloring of the path."""
        self._rainbow_enabled = enabled
        self._rebuild_full_path()

    def toggle_rainbow(self):
        """Toggle rainbow coloring of the path."""
        self._rainbow_enabled = not self._rainbow_enabled
        self._rebuild_full_path()

    def set_savgol_window_length(self, wlen: int):
        """Set the window length for Savitzky-Golay smoothing."""
        wlen = max(3, wlen)
        if wlen % 2 == 0:
            wlen += 1
        self._savgol_window_length = wlen

        self._rebuild_full_path()

    # --------------------------------------------------------------------
    # LOADING
    # --------------------------------------------------------------------
    def load_image(self, path: str):
        """Load an image from a file path."""
        pixmap = QPixmap(path)
        if not pixmap.isNull():
            self.image_item.setPixmap(pixmap)
            self.setSceneRect(QRectF(pixmap.rect()))

            self._img_w = pixmap.width()
            self._img_h = pixmap.height()

            self._clear_all_points()
            self.resetTransform()
            self.fitInView(self.image_item, Qt.KeepAspectRatio)

            # By default, add S/E
            s_x, s_y = 0.15 * self._img_w, 0.5 * self._img_h
            e_x, e_y = 0.85 * self._img_w, 0.5 * self._img_h
            self._insert_anchor_point(-1, s_x, s_y, label="S", removable=False, z_val=100, radius=6)
            self._insert_anchor_point(-1, e_x, e_y, label="E", removable=False, z_val=100, radius=6)

    # --------------------------------------------------------------------
    # ANCHOR POINTS
    # --------------------------------------------------------------------
    def _insert_anchor_point(self, idx, x: float, y: float, label="", removable=True,
                             z_val=0, radius=4):
        """Insert an anchor point at a specific index."""
        x_clamped = self._clamp(x, radius, self._img_w - radius)
        y_clamped = self._clamp(y, radius, self._img_h - radius)

        if idx < 0:
            # Insert before E if there's at least 2 anchors
            if len(self.anchor_points) >= 2:
                idx = len(self.anchor_points) - 1
            else:
                idx = len(self.anchor_points)

        self.anchor_points.insert(idx, (x_clamped, y_clamped))
        color = Qt.green if label in ("S", "E") else Qt.red
        item = LabeledPointItem(x_clamped, y_clamped,
                                label=label, radius=radius, color=color,
                                removable=removable, z_value=z_val)
        self.point_items.insert(idx, item)
        self.scene.addItem(item)

    def _add_guide_point(self, x, y):
        """Add a guide point to the path."""
        x_clamped = self._clamp(x, self.dot_radius, self._img_w - self.dot_radius)
        y_clamped = self._clamp(y, self.dot_radius, self._img_h - self.dot_radius)

        self._revert_cost_to_original()

        if not self._full_path_xy:
            self._insert_anchor_point(-1, x_clamped, y_clamped,
                                      label="", removable=True, z_val=1, radius=self.dot_radius)
        else:
            self._insert_anchor_between_subpath(x_clamped, y_clamped)

        self._apply_all_guide_points_to_cost()
        self._rebuild_full_path()

    def _insert_anchor_between_subpath(self, x_new: float, y_new: float ):
        """Insert an anchor point between existing anchor points."""        # If somehow we have no path yet
        # If somehow we have no path yet
        if not self._full_path_xy:
            self._insert_anchor_point(-1, x_new, y_new)
            return

        # Find nearest point in the current full path
        best_idx = None
        best_d2 = float('inf')
        for i, (px, py) in enumerate(self._full_path_xy):
            dx = px - x_new
            dy = py - y_new
            d2 = dx*dx + dy*dy
            if d2 < best_d2:
                best_d2 = d2
                best_idx = i

        if best_idx is None:
            self._insert_anchor_point(-1, x_new, y_new)
            return

        def approx_equal(xa, ya, xb, yb, tol=1e-3):
            """Check if two points are approximately equal."""
            return (abs(xa - xb) < tol) and (abs(ya - yb) < tol)

        def is_anchor(coord):
            """Check if a point is an anchor point."""
            cx, cy = coord
            for (ax, ay) in self.anchor_points:
                if approx_equal(ax, ay, cx, cy):
                    return True
            return False

        # Walk left
        left_anchor_pt = None
        iL = best_idx
        while iL >= 0:
            px, py = self._full_path_xy[iL]
            if is_anchor((px, py)):
                left_anchor_pt = (px, py)
                break
            iL -= 1

        # Walk right
        right_anchor_pt = None
        iR = best_idx
        while iR < len(self._full_path_xy):
            px, py = self._full_path_xy[iR]
            if is_anchor((px, py)):
                right_anchor_pt = (px, py)
                break
            iR += 1

        # If we can't find distinct anchors on left & right,
        # just insert before E.
        if not left_anchor_pt or not right_anchor_pt:
            self._insert_anchor_point(-1, x_new, y_new)
            return
        if left_anchor_pt == right_anchor_pt:
            self._insert_anchor_point(-1, x_new, y_new)
            return

        # Convert anchor coords -> anchor_points indices
        left_idx = None
        right_idx = None
        for i, (ax, ay) in enumerate(self.anchor_points):
            if approx_equal(ax, ay, left_anchor_pt[0], left_anchor_pt[1]):
                left_idx = i
            if approx_equal(ax, ay, right_anchor_pt[0], right_anchor_pt[1]):
                right_idx = i

        if left_idx is None or right_idx is None:
            self._insert_anchor_point(-1, x_new, y_new)
            return

        # Insert between them
        if left_idx < right_idx:
            insert_idx = left_idx + 1
        else:
            insert_idx = right_idx + 1

        self._insert_anchor_point(insert_idx, x_new, y_new, label="", removable=True,
                                  z_val=1, radius=self.dot_radius)

    # --------------------------------------------------------------------
    # COST IMAGE
    # --------------------------------------------------------------------
    def _revert_cost_to_original(self):
        if self.cost_image_original is not None:
            self.cost_image = self.cost_image_original.copy()

    def _apply_all_guide_points_to_cost(self):
        if self.cost_image is None:
            return
        for i, (ax, ay) in enumerate(self.anchor_points):
            if self.point_items[i].is_removable():
                self._lower_cost_in_circle(ax, ay, self.radius_cost_image)

    def _lower_cost_in_circle(self, x_f: float, y_f: float, radius: int):
        """Lower the cost in a circle centered at (x_f, y_f)."""
        if self.cost_image is None:
            return
        h, w = self.cost_image.shape
        row_c = int(round(y_f))
        col_c = int(round(x_f))
        if not (0 <= row_c < h and 0 <= col_c < w):
            return
        global_min = self.cost_image.min()
        r_s = max(0, row_c - radius)
        r_e = min(h, row_c + radius + 1)
        c_s = max(0, col_c - radius)
        c_e = min(w, col_c + radius + 1)
        for rr in range(r_s, r_e):
            for cc in range(c_s, c_e):
                dist = math.sqrt((rr - row_c)**2 + (cc - col_c)**2)
                if dist <= radius:
                    self.cost_image[rr, cc] = global_min

    # --------------------------------------------------------------------
    # PATH BUILDING
    # --------------------------------------------------------------------
    def _rebuild_full_path(self):
        """Rebuild the full path based on the anchor points."""
        for item in self.full_path_points:
            self.scene.removeItem(item)
        self.full_path_points.clear()
        self._full_path_xy.clear()

        if len(self.anchor_points) < 2 or self.cost_image is None:
            return

        big_xy = []
        for i in range(len(self.anchor_points) - 1):
            xA, yA = self.anchor_points[i]
            xB, yB = self.anchor_points[i + 1]
            sub_xy = self._compute_subpath_xy(xA, yA, xB, yB)
            if i == 0:
                big_xy.extend(sub_xy)
            else:
                if len(sub_xy) > 1:
                    big_xy.extend(sub_xy[1:])

        if len(big_xy) >= self._savgol_window_length:
            arr_xy = np.array(big_xy)
            smoothed = savgol_filter(
                arr_xy,
                window_length=self._savgol_window_length,
                polyorder=2,
                axis=0
            )
            big_xy = smoothed.tolist()

        self._full_path_xy = big_xy[:]

        n_points = len(big_xy)
        for i, (px, py) in enumerate(big_xy):
            fraction = i / (n_points - 1) if n_points > 1 else 0
            color = Qt.red
            if self._rainbow_enabled:
                color = self._rainbow_color(fraction)

            path_item = LabeledPointItem(px, py, label="",
                                         radius=self.path_radius,
                                         color=color,
                                         removable=False,
                                         z_value=0)
            self.full_path_points.append(path_item)
            self.scene.addItem(path_item)

        # Keep anchor labels on top
        for p_item in self.point_items:
            if p_item._text_item:
                p_item.setZValue(100)

    def _compute_subpath_xy(self, xA: float, yA: float, xB: float, yB: float):
        """Compute a subpath between two points."""
        if self.cost_image is None:
            return []
        h, w = self.cost_image.shape
        rA, cA = int(round(yA)), int(round(xA))
        rB, cB = int(round(yB)), int(round(xB))
        rA = max(0, min(rA, h - 1))
        cA = max(0, min(cA, w - 1))
        rB = max(0, min(rB, h - 1))
        cB = max(0, min(cB, w - 1))
        try:
            path_rc = find_path(self.cost_image, [(rA, cA), (rB, cB)])
        except ValueError as e:
            print("Error in find_path:", e)
            return []
        # Convert from (row, col) to (x, y)
        return [(c, r) for (r, c) in path_rc]

    def _rainbow_color(self, fraction: float):
        """Get a rainbow color."""
        hue = int(300 * fraction)
        saturation = 255
        value = 255
        return QColor.fromHsv(hue, saturation, value)

    # --------------------------------------------------------------------
    # MOUSE EVENTS
    # --------------------------------------------------------------------
    def mousePressEvent(self, event):
        """Handle mouse press events for dragging a point or adding a point."""
        if event.button() == Qt.LeftButton:
            self._mouse_pressed = True
            self._was_dragging = False
            self._press_view_pos = event.pos()

            idx = self._find_item_near(event.pos(), threshold=10)
            if idx is not None:
                self._dragging_idx = idx
                self._drag_counter = 0
                scene_pos = self.mapToScene(event.pos())
                px, py = self.point_items[idx].get_pos()
                self._drag_offset = (scene_pos.x() - px, scene_pos.y() - py)
                self.setCursor(Qt.ClosedHandCursor)
                return

        elif event.button() == Qt.RightButton:
            self._remove_point_by_click(event.pos())

        super().mousePressEvent(event)

    def mouseMoveEvent(self, event):
        """Handle mouse move events for dragging a point or dragging the view"""
        if self._dragging_idx is not None:
            scene_pos = self.mapToScene(event.pos())
            x_new = scene_pos.x() - self._drag_offset[0]
            y_new = scene_pos.y() - self._drag_offset[1]

            r = self.point_items[self._dragging_idx]._r
            x_clamped = self._clamp(x_new, r, self._img_w - r)
            y_clamped = self._clamp(y_new, r, self._img_h - r)
            self.point_items[self._dragging_idx].set_pos(x_clamped, y_clamped)

            self._drag_counter += 1
            # Update path every 4 moves
            if self._drag_counter >= 4:
                self._drag_counter = 0
                self._revert_cost_to_original()
                self._apply_all_guide_points_to_cost()
                self.anchor_points[self._dragging_idx] = (x_clamped, y_clamped)
                self._rebuild_full_path()
        else:
            if self._mouse_pressed and (event.buttons() & Qt.LeftButton):
                dist = (event.pos() - self._press_view_pos).manhattanLength()
                if dist > self._drag_threshold:
                    self._was_dragging = True

        super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        """Handle mouse release events for dragging a point or adding a point."""
        super().mouseReleaseEvent(event)
        if event.button() == Qt.LeftButton and self._mouse_pressed:
            self._mouse_pressed = False
            self.setCursor(Qt.ArrowCursor)

            if self._dragging_idx is not None:
                idx = self._dragging_idx
                self._dragging_idx = None
                self._drag_offset = (0, 0)
                newX, newY = self.point_items[idx].get_pos()
                self.anchor_points[idx] = (newX, newY)
                self._revert_cost_to_original()
                self._apply_all_guide_points_to_cost()
                self._rebuild_full_path()
            else:
                # No drag => add point
                if not self._was_dragging:
                    scene_pos = self.mapToScene(event.pos())
                    x, y = scene_pos.x(), scene_pos.y()
                    self._add_guide_point(x, y)

            self._was_dragging = False

    def _remove_point_by_click(self, view_pos: QPoint):
        """Remove a point by clicking on it."""
        idx = self._find_item_near(view_pos, threshold=10)
        if idx is None:
            return
        if not self.point_items[idx].is_removable():
            return

        self.scene.removeItem(self.point_items[idx])
        self.point_items.pop(idx)
        self.anchor_points.pop(idx)

        self._revert_cost_to_original()
        self._apply_all_guide_points_to_cost()
        self._rebuild_full_path()

    def _find_item_near(self, view_pos: QPoint, threshold=10):
        """Find the index of an item near a given position."""
        scene_pos = self.mapToScene(view_pos)
        x_click, y_click = scene_pos.x(), scene_pos.y()

        closest_idx = None
        min_dist = float('inf')
        for i, itm in enumerate(self.point_items):
            d = itm.distance_to(x_click, y_click)
            if d < min_dist:
                min_dist = d
                closest_idx = i
        if closest_idx is not None and min_dist <= threshold:
            return closest_idx
        return None

    # --------------------------------------------------------------------
    # UTILS
    # --------------------------------------------------------------------
    def _clamp(self, val, mn, mx):
        return max(mn, min(val, mx))

    def _clear_all_points(self):
        """Clear all anchor points and guide points."""
        for it in self.point_items:
            self.scene.removeItem(it)
        self.point_items.clear()
        self.anchor_points.clear()

        for p in self.full_path_points:
            self.scene.removeItem(p)
        self.full_path_points.clear()
        self._full_path_xy.clear()

    def clear_guide_points(self):
        """Clear all guide points."""
        i = 0
        while i < len(self.anchor_points):
            if self.point_items[i].is_removable():
                self.scene.removeItem(self.point_items[i])
                del self.point_items[i]
                del self.anchor_points[i]
            else:
                i += 1

        for it in self.full_path_points:
            self.scene.removeItem(it)
        self.full_path_points.clear()
        self._full_path_xy.clear()

        self._revert_cost_to_original()
        self._apply_all_guide_points_to_cost()
        self._rebuild_full_path()

    def get_full_path_xy(self):
        """Returns the entire path as a list of (x, y) coordinates."""
        return self._full_path_xy