import math
from scipy.signal import savgol_filter
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsPixmapItem
from PyQt5.QtGui import QPixmap, QColor
from PyQt5.QtCore import Qt, QRectF
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.
Attributes:
scene (QGraphicsScene): The graphics scene for displaying items.
image_item (QGraphicsPixmapItem): The item for displaying the loaded image.
anchor_points (list): List of tuples representing anchor points (x, y).
point_items (list): List of LabeledPointItem objects for anchor points.
full_path_points (list): List of QGraphicsEllipseItems representing the path.
_full_path_xy (list): List of coordinates for the entire path.
dot_radius (int): Radius of the anchor points.
path_radius (int): Radius of the path points.
radius_cost_image (int): Radius for lowering cost in the cost image.
_img_w (int): Width of the loaded image.
_img_h (int): Height of the loaded image.
_mouse_pressed (bool): Indicates if the mouse is pressed.
_press_view_pos (QPoint): Position of the mouse press event.
_drag_threshold (int): Threshold for detecting drag events.
_was_dragging (bool): Indicates if a drag event occurred.
_dragging_idx (int): Index of the point being dragged.
_drag_offset (tuple): Offset for dragging points.
_drag_counter (int): Counter for drag events.
cost_image_original (np.ndarray): Original cost image.
cost_image (np.ndarray): Current cost image.
_rainbow_enabled (bool): Indicates if rainbow coloring is enabled.
_savgol_window_length (int): Window length for Savitzky-Golay smoothing.
"""
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 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
i_l = best_idx
while i_l >= 0:
px, py = self._full_path_xy[i_l]
if is_anchor((px, py)):
left_anchor_pt = (px, py)
break
i_l -= 1
# Walk right
right_anchor_pt = None
i_r = best_idx
while i_r < len(self._full_path_xy):
px, py = self._full_path_xy[i_r]
if is_anchor((px, py)):
right_anchor_pt = (px, py)
break
i_r += 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 mouse_press_event(self, event):
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().mouse_press_event(event)
def mouse_move_event(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().mouse_move_event(event)
def mouse_release_event(self, event):
super().mouse_release_event(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"):
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):
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):
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