From 548ea30b8452a8aad41c38f2fe3de2356f69d985 Mon Sep 17 00:00:00 2001 From: anna-grim Date: Sun, 12 Jul 2026 19:48:39 +0000 Subject: [PATCH 1/2] refactor: improved utils and curve viz --- .../geometric_learning/curve_visualization.py | 229 +++++++++++++++--- src/neuron_proofreader/skeleton_graph.py | 5 +- src/neuron_proofreader/utils/util.py | 60 ++--- 3 files changed, 212 insertions(+), 82 deletions(-) diff --git a/src/neuron_proofreader/geometric_learning/curve_visualization.py b/src/neuron_proofreader/geometric_learning/curve_visualization.py index 95f7cb8..24db69f 100644 --- a/src/neuron_proofreader/geometric_learning/curve_visualization.py +++ b/src/neuron_proofreader/geometric_learning/curve_visualization.py @@ -21,32 +21,26 @@ # --- Plot Curves --- def plot_curves(curve1, curve2, name1=None, name2=None): + """ + Plots two 3D curves using Plotly. + + Parameters + ---------- + curve1 : numpy.ndarray + Array with shape (N, 3) containing coordinates of the first curve. + curve2 : numpy.ndarray + Array with shape (M, 3) containing coordinates of the second curve. + name1 : str, optional + Label for the first curve in the plot legend. Default is None. + name2 : str, optional + Label for the second curve in the plot legend. Default is None. + """ + pt = np.zeros((1, 3)) fig = go.Figure( data=[ - go.Scatter3d( - x=curve1[:, 0], - y=curve1[:, 1], - z=curve1[:, 2], - mode="lines", - name=name1, - line=dict(width=5, color="blue"), - ), - go.Scatter3d( - x=curve2[:, 0], - y=curve2[:, 1], - z=curve2[:, 2], - mode="lines", - name=name2, - line=dict(width=5, color="green"), - ), - go.Scatter3d( - x=[0], - y=[0], - z=[0], - mode="markers", - name="Origin", - marker=dict(size=3, color="red"), - ), + create_scatter3d(curve1, color="blue", name=name1), + create_scatter3d(curve2, color="green", name=name2), + create_scatter3d(pt, color="red", mode="markers", name="Origin"), ] ) fig.update_layout( @@ -55,18 +49,31 @@ def plot_curves(curve1, curve2, name1=None, name2=None): fig.show() -def plot_length_distribution(dataset_collection, title=None, output_path=None): +def plot_length_distribution(dataset, output_path=None, title=None): + """ + Plots the distribution of path lengths in a dataset. + + Parameters + ---------- + dataset : PathsDataset + Dataset containing an "examples_df" attribute with a "length" column + specifying the path lengths. + output_path : str, optional + If provided, the figure is saved to this location. Otherwise, it is + displayed. Default is None. + title : str, optional + Title of the plot. Default is None. + """ # Compute path length stats - lengths = dataset_collection.examples_df["length"] - p50 = round(np.percentile(lengths, 50), 2) - p99 = round(np.percentile(lengths, 99.9), 2) - thr_lengths = [l for l in lengths if l < p99] + lengths = dataset.examples_df["length"].to_numpy() + p50, p99 = np.percentile(lengths, [50, 99.9]) + lengths = lengths[lengths <= p999] # Plot path lengths plt.figure(figsize=(8, 5)) - plt.hist(thr_lengths, bins=50, edgecolor="white", linewidth=0.5, zorder=2) - add_line(p50, color="r", label=f"50th perc = {p50}") - add_line(p99, color="g", label=f"99.9th perc = {p99}") + plt.hist(lengths, bins=50, edgecolor="white", linewidth=0.5, zorder=2) + add_line(p50, color="r", label=f"50th perc = {p50:.2f}") + add_line(p99, color="g", label=f"99.9th perc = {p99:.2f}") # Plot labels plt.grid(axis="y", color="lightgrey", linewidth=0.5) @@ -90,13 +97,26 @@ def plot_length_distribution(dataset_collection, title=None, output_path=None): # --- Plot Curve Embeddings --- def plot_error_vs_length(lengths, rmse_results, output_path=None): + """ + Plots reconstruction error as a function of path length. + + ---------- + lengths : numpy.ndarray + One-dimensional array of path lengths in mircons. + rmse_results : numpy.ndarray + One-dimensional array of root mean squared errors in mircons. + output_path : str, optional + If provided, the figure is saved to this location. Otherwise, it is + displayed. Default is None. + """ # Set colors norm = LogNorm(vmin=lengths.min(), vmax=lengths.max()) colors = plt.cm.viridis(norm(lengths)) # Plot plt.figure(figsize=(8, 6)) - plt.scatter(lengths, rmse_results, c=colors, s=10, alpha=0.8) + sc = plt.scatter(lengths, rmse_results, c=colors, s=10, alpha=0.8) + plt.colorbar(sc, label="Path Length (μm)") plt.xscale("log") plt.yscale("log") plt.xlabel("Path Length (μm)", fontsize=12) @@ -105,6 +125,19 @@ def plot_error_vs_length(lengths, rmse_results, output_path=None): def plot_latents_by_pca(curves, latents, output_path=None): + """ + Visualizes latent representations using PCA. + + Parameters + ---------- + curves : List[numpy.ndarray] + Curves such that each is an array with shape (N, 3). + latents : numpy.ndarray + Array of latent representations with shape (N, latent_dim). + output_path : str, optional + Base path for saving the figures. If provided, the figure is saved to + this location. Otherwise, it is displayed. Default is None. + """ # Set output paths if output_path: dir_output_path = output_path.replace("pca", "pca_direction") @@ -124,6 +157,21 @@ def plot_latents_by_pca(curves, latents, output_path=None): def _plot_latents_by_direction(curves, latents_2d, pca, output_path=None): + """ + Plots 2D latent colored by the principal direction of each curve. + + Parameters + ---------- + curves : List[numpy.ndarray] + Curves such that each is an array with shape (N, 3). + latents_2d : numpy.ndarray + Latent representations with shape (N, 2). + pca : sklearn.decomposition.PCA + Fitted PCA object used to generate "latents_2d". + output_path : str, optional + If provided, the figure is saved to this location. Otherwise, it is + displayed. Default is None. + """ # Compute directions and colors for each curve directions = np.array([curve_principal_direction(c) for c in curves]) colors = np.array([direction_to_color(d) for d in directions]) @@ -138,6 +186,21 @@ def _plot_latents_by_direction(curves, latents_2d, pca, output_path=None): def _plot_latents_by_length(lengths, latents_2d, pca, output_path=None): + """ + Plots 2D latents colored by the length of each curve. + + Parameters + ---------- + curves : List[numpy.ndarray] + Curves such that each is an array with shape (N, 3). + latents_2d : numpy.ndarray + Latent representations with shape (N, 2). + pca : sklearn.decomposition.PCA + Fitted PCA object used to generate "latents_2d". + output_path : str, optional + If provided, the figure is saved to this location. Otherwise, it is + displayed. Default is None. + """ plt.figure(figsize=(8, 6)) sc = plt.scatter( latents_2d[:, 0], @@ -148,19 +211,87 @@ def _plot_latents_by_length(lengths, latents_2d, pca, output_path=None): alpha=0.7, norm=LogNorm(vmin=lengths.min(), vmax=lengths.max()), ) - plt.colorbar(sc, label="Path length (μm)") + plt.colorbar(sc, label="Path Length (μm)") plt.xlabel(f"PC1 ({pca.explained_variance_ratio_[0]*100:.1f}%)") plt.ylabel(f"PC2 ({pca.explained_variance_ratio_[1]*100:.1f}%)") - plt.title("PCA of curve embeddings") + plt.title("Curve Embeddings Colored by Path Length") visualize_result(output_path=output_path) # --- Helpers --- -def add_line(p, color=None, label=None): - plt.axvline(p, color=color, linestyle="--", label=label, zorder=2) +def add_line(x, color=None, label=None): + """ + Adds a vertical reference line to the current Matplotlib axes. + + Parameters + ---------- + x : float + X-coordinate at which to draw the vertical line. + color : str, optional + Color of the line. Default is None. + label : str, optional + Label for the line displayed in the plot legend. Default is None. + """ + plt.axvline(x, color=color, linestyle="--", label=label, zorder=2) + + +def create_scatter3d(pts, color=None, mode="lines", name=None, width=5): + """ + Creates a Plotly 3D scatter trace. + + Parameters + ---------- + pts : numpy.ndarray + Array with shape (N, 3) containing the 3D coordinates to plot. + color : str, optional + Color of the line or markers. Default is None. + mode : str, optional + Rendering mode for the trace. Default is "lines". + name : str, optional + Name of the trace displayed in the plot legend. Default is None. + width : float, optional + Width of object to be plotted. Default is 5. + + Returns + ------- + plotly.graph_objects.Scatter3d + A Plotly 3D scatter trace. + """ + # Create object dict + if mode == "lines": + line = dict(width=width, color=color) + marker = None + else: + line = None + marker = dict(size=width, color=color) + + # Create scatter plot + return go.Scatter3d( + x=pts[:, 0], + y=pts[:, 1], + z=pts[:, 2], + mode=mode, + name=name, + line=line, + marker=marker, + ) 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] @@ -170,6 +301,21 @@ def curve_principal_direction(curve): def direction_to_color(direction): + """ + Converts a 3D direction vector into an RGB color representation, where the + azimuth angle of the vector is mapped to a hue and polar angle is mapped + to a saturation. + + Parameters + ---------- + direction : numpy.ndarray + Unit vector of shape (3,) representing a 3D direction. + + Returns + ------- + tuple + RGB color corresponding to the input direction. + """ x, y, z = direction azimuth = np.arctan2(y, x) hue = (azimuth + np.pi) / (2 * np.pi) @@ -180,6 +326,15 @@ def direction_to_color(direction): def visualize_result(output_path=None): + """ + Displays or saves the current Matplotlib figure. + + Parameters + ---------- + output_path : str, optional + If provided, the figure is saved to this location. Otherwise, it is + displayed. Default is None. + """ plt.tight_layout() if output_path: plt.savefig(output_path, dpi=300, bbox_inches="tight") diff --git a/src/neuron_proofreader/skeleton_graph.py b/src/neuron_proofreader/skeleton_graph.py index 1bd13a9..73c5226 100644 --- a/src/neuron_proofreader/skeleton_graph.py +++ b/src/neuron_proofreader/skeleton_graph.py @@ -626,10 +626,9 @@ def clip_to_bbox(self, metadata_path): Path to JSON file containing origin and shape of bounding box to clip skeletons to. """ - bucket_name, path = util.parse_cloud_path(metadata_path) - if util.check_gcs_file_exists(bucket_name, path): + if util.check_gcs_file_exists(metadata_path): # Extract bounding box - metadata = util.read_json_from_gcs(bucket_name, path) + metadata = util.read_json(metadata_path) origin = metadata["chunk_origin"][::-1] shape = metadata["chunk_shape"][::-1] diff --git a/src/neuron_proofreader/utils/util.py b/src/neuron_proofreader/utils/util.py index 31c65ca..0684273 100644 --- a/src/neuron_proofreader/utils/util.py +++ b/src/neuron_proofreader/utils/util.py @@ -217,8 +217,11 @@ def read_json(path): dict Contents of JSON file. """ - with open(path, "r") as f: - return json.load(f) + if is_gcs_path(path): + return json.loads(read_gcs_txt(path)) + else: + with open(path, "r") as f: + return json.load(f) def read_txt(path, client=None): @@ -338,9 +341,9 @@ def get_google_swcs_prefix(root_prefix, brain_id, segmentation_id): # Determine old vs. new result prefix1 = os.path.join(root_prefix, brain_id, "whole_brain") prefix2 = os.path.join(root_prefix, "whole_brain", brain_id) - if check_gcs_exists(prefix1, is_prefix=True): + if check_gcs_prefix_exists(prefix1): prefix = prefix1 - elif check_gcs_exists(prefix2, is_prefix=True): + elif check_gcs_prefix_exists(prefix2): prefix = prefix2 else: raise Exception("Unable to find Google swcs result!") @@ -406,21 +409,20 @@ def parse_cloud_path(path): # Split path parts = path.split("/", 1) bucket_name = parts[0] - prefix = parts[1] if len(parts) > 1 else "" - return bucket_name, prefix + key = parts[1] if len(parts) > 1 else "" + return bucket_name, key # --- GCS Utils --- -def check_gcs_exists(path, is_prefix=False): +def check_gcs_file_exists(path): """ - Checks if a file or prefix exists in GCS. + Checks if a file exists at the given GCS path. + Parameters ---------- path : str GCS path to check. - prefix : bool - If True, checks whether any object exists under the given prefix. - If False, checks whether the exact file exists. + Returns ------- bool @@ -428,18 +430,14 @@ def check_gcs_exists(path, is_prefix=False): """ bucket_name, key = parse_cloud_path(path) bucket = storage.Client().bucket(bucket_name) - if is_prefix: - key = key.rstrip("/") + "/" - return any(bucket.list_blobs(prefix=key, max_results=1)) - else: - return bucket.blob(key).exists() + return bucket.blob(key).exists() -def check_gcs_prefix_exists(path): - bucket_name, prefix = parse_cloud_path(path) - prefix = prefix.rstrip("/") + "/" +def check_gcs_prefix_exists(prefix): + bucket_name, key = parse_cloud_path(prefix) bucket = storage.Client().bucket(bucket_name) - exists = any(bucket.list_blobs(prefix=prefix, max_results=1)) + key = key.rstrip("/") + "/" + exists = any(bucket.list_blobs(prefix=key, max_results=1)) return exists @@ -540,28 +538,6 @@ def read_gcs_txt(prefix, client=None): return bucket.blob(subprefix).download_as_text() -def read_json_from_gcs(bucket_name, blob_path): - """ - Reads JSON file stored in a GCS bucket. - - Parameters - ---------- - bucket_name : str - Name of the GCS bucket containing the JSON file. - blob_path : str - Path to the JSON file within the GCS bucket. - - Returns - ------- - dict - Parsed JSON content as a Python dictionary. - """ - client = storage.Client() - bucket = client.bucket(bucket_name) - blob = bucket.blob(blob_path) - return json.loads(blob.download_as_text()) - - # --- S3 Utils --- def is_s3_path(path): """ From 4561cb799e457241ea3232d5d588a885952f8c1b Mon Sep 17 00:00:00 2001 From: anna-grim Date: Sun, 12 Jul 2026 20:02:24 +0000 Subject: [PATCH 2/2] bug: plot path lengths --- .../geometric_learning/curve_visualization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/neuron_proofreader/geometric_learning/curve_visualization.py b/src/neuron_proofreader/geometric_learning/curve_visualization.py index 24db69f..116df54 100644 --- a/src/neuron_proofreader/geometric_learning/curve_visualization.py +++ b/src/neuron_proofreader/geometric_learning/curve_visualization.py @@ -67,7 +67,7 @@ def plot_length_distribution(dataset, output_path=None, title=None): # Compute path length stats lengths = dataset.examples_df["length"].to_numpy() p50, p99 = np.percentile(lengths, [50, 99.9]) - lengths = lengths[lengths <= p999] + lengths = lengths[lengths <= p99] # Plot path lengths plt.figure(figsize=(8, 5))