diff --git a/src/neuron_proofreader/geometric_learning/curve_visualization.py b/src/neuron_proofreader/geometric_learning/curve_visualization.py index 220d12c..be8e81c 100644 --- a/src/neuron_proofreader/geometric_learning/curve_visualization.py +++ b/src/neuron_proofreader/geometric_learning/curve_visualization.py @@ -12,12 +12,18 @@ from matplotlib.colors import LogNorm from sklearn.decomposition import PCA +import matplotlib.patheffects as pe import matplotlib.pyplot as plt import numpy as np import plotly.graph_objects as go from neuron_proofreader.utils import geometry_util +CURVE_COLORS = [ + "#0072B2", # blue + "#D55E00", # vermillion +] + # --- Plot Curves --- def plot_curves(curve1, curve2, name1=None, name2=None): @@ -38,8 +44,8 @@ def plot_curves(curve1, curve2, name1=None, name2=None): pt = np.zeros((1, 3)) fig = go.Figure( data=[ - create_scatter3d(curve1, color="blue", name=name1), - create_scatter3d(curve2, color="green", name=name2), + create_scatter3d(curve1, color=CURVE_COLORS[0], name=name1), + create_scatter3d(curve2, color=CURVE_COLORS[1], name=name2), create_scatter3d(pt, color="red", mode="markers", name="Origin"), ] ) @@ -95,6 +101,140 @@ def plot_length_distribution(dataset, output_path=None, title=None): visualize_result(output_path=output_path) +# --- Plot Publication Ready Curves --- +def plot_curve_mip( + curve1, + curve2, + name1="Original", + name2="Reconstruction", + output_path=None, +): + """ + Creates a publication-quality PCA projection of two 3D curves. + + Parameters + ---------- + curve1 : numpy.ndarray + First 3D curve. + curve2 : numpy.ndarray + Second 3D curve. + name1 : str + Label for first curve. + name2 : str + Label for second curve. + output_path : str, optional + Output path. + """ + # Project 3D curves + curve1, curve2 = geometry_util.curves_pca_projection(curve1, curve2) + pts = np.vstack([curve1, curve2]) + + # Create figure + fig, ax = plt.subplots(figsize=(7, 5), dpi=300) + fig.subplots_adjust( + left=0.03, + right=0.97, + bottom=0.08, + top=0.95, + ) + + # Plot curves + itr = zip((curve1, curve2), CURVE_COLORS, (name1, name2)) + for curve, color, label in itr: + ax.plot( + curve[:, 0], + curve[:, 1], + color=color, + linewidth=3, + solid_capstyle="round", + alpha=0.9, + label=label, + ) + + # Preserve geometry but allow rectangular neurons + ax.set_aspect( + "equal", + adjustable="datalim", + ) + ax.autoscale() + ax.margins(0.05) + ax.axis("off") + + # Add annotations + scale_length = nice_scale_bar_length(np.ptp(pts[:, 0]) * 0.12) + _add_scale_bar(fig, ax, scale_length) + _add_error_annotation(fig, ax, curve1, curve2) + + ax.legend( + loc="upper left", + fontsize=9, + frameon=True, + framealpha=0.8, + edgecolor="none", + ) + + visualize_result(output_path=output_path) + + +def _add_scale_bar(fig, ax, length, width=0.1): + """ + Adds a scale bar in the bottom-left of the figure. + """ + # Get scale bar position + x0 = ax.get_position().x0 + 0.02 + y0 = ax.get_position().y0 + 0.05 + + # Create scale bar + bar = plt.Line2D( + [x0, x0 + width], + [y0, y0], + transform=fig.transFigure, + color="black", + linewidth=4, + solid_capstyle="butt", + ) + + path_effects = [pe.Stroke(linewidth=7, foreground="white"), pe.Normal()] + bar.set_path_effects(path_effects) + + # Add text + fig.lines.append(bar) + fig.text( + x0 + width / 2, + y0 + 0.015, + f"{length:g} μm", + transform=fig.transFigure, + ha="center", + va="bottom", + fontsize=9, + ) + + +def _add_error_annotation(fig, ax, curve1, curve2): + """ + Adds path length and reconstruction error annotations. + """ + # Compute error + gt_length = geometry_util.path_length(curve1) + recon_length = geometry_util.path_length(curve2) + error = geometry_util.max_l2_error(curve1, curve2) + + # Display reconstruction error + fig.text( + ax.get_position().x1 - 0.02, + ax.get_position().y0 + 0.05, + ( + f"GT length: {gt_length:.1f}μm\n" + f"Recon length: {recon_length:.1f} μm\n" + f"Max L2 error: {error:.2f} μm" + ), + transform=fig.transFigure, + fontsize=9, + ha="right", + va="bottom", + ) + + # --- Plot Curve Embeddings --- def plot_error_vs_length(lengths, rmse_results, output_path=None): """ @@ -173,7 +313,7 @@ def _plot_latents_by_direction(curves, latents_2d, pca, output_path=None): displayed. Default is None. """ # Compute directions and colors for each curve - directions = np.array([curve_principal_direction(c) for c in curves]) + directions = [geometry_util.curve_principal_direction(c) for c in curves] colors = np.array([direction_to_color(d) for d in directions]) # Plot @@ -277,29 +417,6 @@ def create_scatter3d(pts, color=None, mode="lines", name=None, width=5): ) -def curve_principal_direction(curve): - """ - Computes the principal direction of a 3D curve using PCA. - - Parameters - ---------- - curve : numpy.ndarray - Array with shape (N, 3) containing the 3D coordinates of the curve. - - Returns - ------- - numpy.ndarray - Unit vector of shape (3,) representing the principal direction of the - curve. - """ - curve_pca = PCA(n_components=1) - curve_pca.fit(curve) - direction = curve_pca.components_[0] - if direction[2] < 0: - direction = -direction - return direction / np.linalg.norm(direction) - - def direction_to_color(direction): """ Converts a 3D direction vector into an RGB color representation, where the @@ -325,6 +442,27 @@ def direction_to_color(direction): return hsv_to_rgb(hue, saturation, value) +def nice_scale_bar_length(length): + """ + Rounds a length to nice numerical value. + + Parameters + ---------- + length : float + Proposed scale bar length. + + Returns + ------- + float + Rounded scale bar length. + """ + exponent = np.floor(np.log10(length)) + fraction = length / 10**exponent + for value in (1, 2, 5, 10): + if fraction < 1.5 * value: + return value * 10**exponent + + def visualize_result(output_path=None): """ Displays or saves the current Matplotlib figure. diff --git a/src/neuron_proofreader/utils/geometry_util.py b/src/neuron_proofreader/utils/geometry_util.py index ce9e256..39866b2 100644 --- a/src/neuron_proofreader/utils/geometry_util.py +++ b/src/neuron_proofreader/utils/geometry_util.py @@ -12,6 +12,7 @@ from scipy.interpolate import UnivariateSpline from scipy.linalg import svd from scipy.spatial.distance import euclidean +from sklearn.decomposition import PCA from tqdm import tqdm import networkx as nx @@ -20,41 +21,76 @@ # --- Curve Utils --- -def path_length(curve): +def max_l2_error(curve1, curve2): """ - Computes the Euclidean length of the given curve. + Computes maximum pointwise L2 error. Parameters ---------- - curve : numpy.ndarray - Array of points that form an n-d curve. + curve1 : numpy.ndarray + Ground truth curve. + curve2 : numpy.ndarray + Reconstruction curve. Returns ------- float - Euclidean length of the given curve. + Maximum Euclidean error. """ - return np.linalg.norm(np.diff(curve, axis=0), axis=1,).sum() + assert curve1.shape == curve2.shape, "Curves have different number of pts" + return np.linalg.norm(curve1 - curve2, axis=1).max() -def compute_max_l2_error(curve1, curve2): +def curves_pca_projection(curve1, curve2): """ - Computes maximum pointwise L2 error. + Projects two 3D curves into a shared PCA coordinate system. Parameters ---------- curve1 : numpy.ndarray - Ground truth curve. + First curve with shape ``(N, 3)``. curve2 : numpy.ndarray - Reconstruction curve. + Second curve with shape ``(M, 3)``. Returns ------- - float - Maximum Euclidean error. + tuple + Two projected curves with shape ``(N, 2)`` and ``(M, 2)``. """ - assert curve1.shape == curve2.shape, "Curves have different number of pts" - return np.linalg.norm(curve1 - curve2, axis=1,).max() + # Compute PCA components + pts = np.vstack([curve1, curve2]) + center = pts.mean(axis=0) + + pca = PCA(n_components=3) + pca.fit(pts - center) + + # Compute projections + curve1_proj = pca.transform(curve1 - center) + curve2_proj = pca.transform(curve2 - center) + return curve1_proj[:, :2], curve2_proj[:, :2] + + +def curve_principal_direction(curve): + """ + Computes the principal direction of a 3D curve using PCA. + + Parameters + ---------- + curve : numpy.ndarray + Array with shape (N, 3) containing the 3D coordinates of the curve. + + Returns + ------- + numpy.ndarray + Unit vector of shape (3,) representing the principal direction of the + curve. + """ + curve_pca = PCA(n_components=1) + curve_pca.fit(curve) + direction = curve_pca.components_[0] + if direction[2] < 0: + direction = -direction + return direction / np.linalg.norm(direction) def fit_spline_1d(pts, k=3, s=None): @@ -108,6 +144,23 @@ def fit_spline_3d(pts, k=3, s=None): return spline_x, spline_y, spline_z +def path_length(curve): + """ + Computes the Euclidean length of the given curve. + + Parameters + ---------- + curve : numpy.ndarray + Array of points that form an n-d curve. + + Returns + ------- + float + Euclidean length of the given curve. + """ + return np.linalg.norm(np.diff(curve, axis=0), axis=1).sum() + + def resample_curve_1d(pts, n_pts=None, s=None): """ Smooths a 1D curve by fitting a spline and resampling it. @@ -350,8 +403,7 @@ def compute_svd(pts): Unitary matrix having right singular vectors as rows. Of shape (D, D) or (K, D) depending on full_matrices. """ - pts = pts - np.mean(pts, axis=0) - return svd(pts) + return svd(pts - np.mean(pts, axis=0)) def make_digital_line(p1, p2):