render_neuron_component_overlays.py

#!/usr/bin/env python3
"""Render per-neuron GT and predicted-component overlay figures."""

from __future__ import annotations

import argparse
import colorsys
import csv
import hashlib
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Sequence, Tuple

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import LineCollection


@dataclass(frozen=True)
class SwcData:
    coords: np.ndarray
    segments: np.ndarray
    length_um: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results-dir", required=True, type=Path)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--lengths-csv", type=Path)
    parser.add_argument("--voxel-size-xy", type=float, default=0.748)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--dpi", type=int, default=220)
    parser.add_argument("--figure-width", type=float, default=14.0)
    parser.add_argument("--figure-height", type=float, default=8.8)
    return parser.parse_args()


def parse_swc(path: Path, plot_xy_scale: float = 1.0, length_xy_scale: float = 1.0) -> SwcData:
    node_ids: List[int] = []
    coords_list: List[Tuple[float, float, float]] = []
    parent_ids: List[int] = []
    with path.open("r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            if len(parts) < 7:
                continue
            try:
                node_ids.append(int(float(parts[0])))
                coords_list.append((float(parts[2]), float(parts[3]), float(parts[4])))
                parent_ids.append(int(float(parts[6])))
            except ValueError:
                continue

    if not node_ids:
        return SwcData(
            coords=np.zeros((0, 3), dtype=float),
            segments=np.zeros((0, 2, 3), dtype=float),
            length_um=0.0,
        )

    raw_coords = np.asarray(coords_list, dtype=float)
    coords = raw_coords.copy()
    coords[:, 0] *= plot_xy_scale
    coords[:, 1] *= plot_xy_scale
    length_coords = raw_coords.copy()
    length_coords[:, 0] *= length_xy_scale
    length_coords[:, 1] *= length_xy_scale
    node_index = {nid: idx for idx, nid in enumerate(node_ids)}
    segments: List[np.ndarray] = []
    length_um = 0.0
    for child_idx, parent_id in enumerate(parent_ids):
        parent_idx = node_index.get(parent_id)
        if parent_idx is None:
            continue
        seg = np.asarray([coords[parent_idx], coords[child_idx]], dtype=float)
        segments.append(seg)
        length_um += float(np.linalg.norm(length_coords[child_idx] - length_coords[parent_idx]))
    return SwcData(
        coords=coords,
        segments=np.asarray(segments, dtype=float) if segments else np.zeros((0, 2, 3), dtype=float),
        length_um=length_um,
    )


def read_length_rows(path: Path | None) -> Dict[str, dict]:
    if path is None or not path.exists():
        return {}
    with path.open("r", newline="") as f:
        rows = list(csv.DictReader(f))
    return {row["Neuron"]: row for row in rows if row.get("Neuron")}


def stable_color(name: str, seed: int) -> Tuple[float, float, float]:
    digest = hashlib.md5(f"{seed}:{name}".encode("utf-8")).hexdigest()
    hue = (int(digest[:8], 16) % 360) / 360.0
    sat = 0.55 + (int(digest[8:12], 16) % 20) / 100.0
    val = 0.72 + (int(digest[12:16], 16) % 18) / 100.0
    return colorsys.hsv_to_rgb(hue, min(sat, 0.85), min(val, 0.95))


def choose_scalebar_length(width_um: float) -> int:
    for length in (5000, 2000, 1000, 500, 200, 100, 50):
        if length <= width_um * 0.28:
            return length
    return max(10, int(width_um * 0.2))


def format_mm(value_mm: float) -> str:
    return f"{value_mm:.2f} mm"


def set_panel_style(ax: plt.Axes) -> None:
    ax.set_facecolor("white")
    ax.set_aspect("equal")
    ax.grid(False)
    ax.invert_yaxis()
    ax.set_xticks([])
    ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set_visible(False)


def add_segments(ax: plt.Axes, segments: np.ndarray, color, linewidth: float, alpha: float) -> None:
    if segments.size == 0:
        return
    ax.add_collection(LineCollection(segments[:, :, :2], colors=[color], linewidths=linewidth, alpha=alpha))


def add_scalebar(
    ax: plt.Axes,
    xlim: Tuple[float, float],
    ylim: Tuple[float, float],
    voxel_size_xy: float,
) -> None:
    width = xlim[1] - xlim[0]
    height = ylim[1] - ylim[0]
    bar_len_um = choose_scalebar_length(width * voxel_size_xy)
    bar_len_plot = bar_len_um / voxel_size_xy
    x0 = xlim[0] + 0.055 * width
    y0 = ylim[1] - 0.08 * height
    ax.plot([x0, x0 + bar_len_plot], [y0, y0], color="black", linewidth=2.3, solid_capstyle="butt")
    ax.text(
        x0 + bar_len_plot / 2.0,
        y0 - 0.018 * height,
        f"{bar_len_um:.0f} um",
        ha="center",
        va="top",
        fontsize=9,
        fontweight="bold",
        color="black",
    )


def panel_limits(*coord_arrays: np.ndarray) -> Tuple[Tuple[float, float], Tuple[float, float]]:
    coords = [arr[:, :2] for arr in coord_arrays if arr.size]
    if not coords:
        return (-1.0, 1.0), (-1.0, 1.0)
    stacked = np.vstack(coords)
    mins = np.min(stacked, axis=0)
    maxs = np.max(stacked, axis=0)
    span = np.maximum(maxs - mins, 1.0)
    pad = np.maximum(span * 0.06, 60.0)
    xlim = (float(mins[0] - pad[0]), float(maxs[0] + pad[0]))
    ylim = (float(mins[1] - pad[1]), float(maxs[1] + pad[1]))
    return xlim, ylim


def render_neuron(
    neuron_id: str,
    gt: SwcData,
    component_items: Sequence[Tuple[str, SwcData]],
    output_path: Path,
    length_row: dict | None,
    seed: int,
    voxel_size_xy: float,
    figure_size: Tuple[float, float],
    dpi: int,
) -> dict:
    pred_coords = [item.coords for _, item in component_items if item.coords.size]
    xlim, ylim = panel_limits(gt.coords, *pred_coords)
    fig = plt.figure(figsize=figure_size, facecolor="white")
    gs = fig.add_gridspec(2, 2, height_ratios=[11.0, 1.7], hspace=0.02, wspace=0.08)
    ax_gt = fig.add_subplot(gs[0, 0])
    ax_pred = fig.add_subplot(gs[0, 1])
    ax_gt_text = fig.add_subplot(gs[1, 0])
    ax_pred_text = fig.add_subplot(gs[1, 1])

    for ax in (ax_gt, ax_pred):
        set_panel_style(ax)
        ax.set_xlim(*xlim)
        ax.set_ylim(*ylim)

    add_segments(ax_gt, gt.segments, color="#6c6c6c", linewidth=2.0, alpha=0.98)
    add_scalebar(ax_gt, xlim, ylim, voxel_size_xy)
    ax_gt.set_title("GT", fontsize=12, fontweight="bold", pad=12)

    add_segments(ax_pred, gt.segments, color="#d0d0d0", linewidth=1.5, alpha=0.9)
    for comp_name, comp_data in component_items:
        if comp_data.segments.size == 0:
            continue
        color = stable_color(comp_name, seed)
        ax_pred.add_collection(
            LineCollection(comp_data.segments[:, :, :2], colors=[color], linewidths=1.9, alpha=0.95)
        )
    add_scalebar(ax_pred, xlim, ylim, voxel_size_xy)
    ax_pred.set_title("Predicted Components", fontsize=12, fontweight="bold", pad=12)

    for ax in (ax_gt_text, ax_pred_text):
        ax.set_axis_off()
        ax.set_facecolor("white")

    gt_mm = float(length_row["GT Length (mm)"]) if length_row and length_row.get("GT Length (mm)") else gt.length_um / 1000.0
    pred_mm = (
        float(length_row["Predicted Components Length (mm)"])
        if length_row and length_row.get("Predicted Components Length (mm)")
        else sum(comp.length_um for _, comp in component_items) / 1000.0
    )
    ax_gt_text.text(0.5, 0.58, format_mm(gt_mm), ha="center", va="center", fontsize=22, color="black")
    ax_pred_text.text(0.5, 0.58, format_mm(pred_mm), ha="center", va="center", fontsize=22, color="black")
    ax_pred_text.text(
        0.5,
        0.12,
        f"{len(component_items)} components",
        ha="center",
        va="center",
        fontsize=10,
        color="#555555",
    )

    fig.suptitle(neuron_id, fontsize=16, fontweight="bold", y=0.985)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    return {
        "neuron_id": neuron_id,
        "gt_length_mm": gt_mm,
        "predicted_components_length_mm": pred_mm,
        "component_count": len(component_items),
        "figure_path": str(output_path),
    }


def main() -> None:
    args = parse_args()
    results_dir = args.results_dir
    output_dir = args.output_dir or (results_dir / "neuron-component-overlays")
    lengths_csv = args.lengths_csv or (results_dir / "merged-length.csv")
    gt_dir = results_dir / "input-swcs-flattened"
    pred_dir = results_dir / "predicted-components"

    length_rows = read_length_rows(lengths_csv)
    summary_rows: List[dict] = []
    figure_size = (args.figure_width, args.figure_height)

    for gt_path in sorted(gt_dir.glob("*.swc")):
        neuron_id = gt_path.stem
        gt = parse_swc(gt_path, plot_xy_scale=1.0, length_xy_scale=args.voxel_size_xy)
        component_dir = pred_dir / neuron_id
        component_items = [
            (
                comp_path.stem,
                parse_swc(
                    comp_path,
                    plot_xy_scale=1.0 / args.voxel_size_xy,
                    length_xy_scale=1.0,
                ),
            )
            for comp_path in sorted(component_dir.glob("*.swc"))
        ]
        out_path = output_dir / f"{neuron_id}_component_overlay.png"
        summary_rows.append(
            render_neuron(
                neuron_id=neuron_id,
                gt=gt,
                component_items=component_items,
                output_path=out_path,
                length_row=length_rows.get(neuron_id),
                seed=args.seed,
                voxel_size_xy=args.voxel_size_xy,
                figure_size=figure_size,
                dpi=args.dpi,
            )
        )

    summary_path = output_dir / "overlay_summary.csv"
    with summary_path.open("w", newline="") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "neuron_id",
                "gt_length_mm",
                "predicted_components_length_mm",
                "component_count",
                "figure_path",
            ],
        )
        writer.writeheader()
        for row in summary_rows:
            writer.writerow(row)


if __name__ == "__main__":
    main()