app.py

"""
Cochlear Neurofilament Tracer — HuggingFace Gradio app
======================================================

Traces auditory-nerve fibers (Neurofilament channel) in confocal z-stacks of
the organ of Corti, uses the Myo7a hair-cell channel to separate
IHC-innervating from OHC-innervating fibers, and reports per-region
quantification (number of fibers, diameter, length, branch points, area
covered) plus a black-background skeleton image and an Excel workbook.

Accepts Zeiss **CZI** z-stacks and generic **TIFF** stacks.
"""

import os
import tempfile
import zipfile
import traceback

import numpy as np
import pandas as pd
import gradio as gr
from skimage.io import imsave
from scipy import ndimage as ndi

import processing as P

OUT_DIR = tempfile.mkdtemp(prefix="neuron_tracer_")

METRIC_COLS = [
    "File", "Frequency region", "Region", "Number of fibers",
    "Hair cells (Myo7a)", "Total length (um)", "Mean diameter (um)",
    "Median diameter (um)", "Branch points", "Area covered (um^2)",
    "FOV area (um^2)", "Area covered (% of FOV)",
]


# --------------------------------------------------------------------------- #
#  Small helpers
# --------------------------------------------------------------------------- #

def _channel_choices(img: P.LoadedImage):
    choices = []
    for i, ch in enumerate(img.channels):
        dye = ch.get("dye") or "no dye / transmitted"
        choices.append((f"Ch {i} — {ch.get('name','?')} ({dye})", i))
    return choices


def _draw_boundary_on(gray_u8, boundary_frac, axis="Y"):
    """Return an RGB copy of a grayscale MIP with a yellow boundary line."""
    rgb = np.stack([gray_u8] * 3, axis=-1)
    ny, nx = gray_u8.shape
    if axis.upper() == "Y":
        b = min(max(int(round(boundary_frac * ny)), 0), ny - 1)
        rgb[b, :] = (255, 255, 0)
    else:
        b = min(max(int(round(boundary_frac * nx)), 0), nx - 1)
        rgb[:, b] = (255, 255, 0)
    return rgb


def _save_png(arr, name):
    path = os.path.join(OUT_DIR, name)
    imsave(path, arr)
    return path


def _build_excel(rows, path):
    """Write a tidy per-region sheet plus a frequency×region summary sheet."""
    df = pd.DataFrame(rows, columns=METRIC_COLS)
    with pd.ExcelWriter(path, engine="openpyxl") as xl:
        df.to_excel(xl, sheet_name="Per region", index=False)
        # Summary: only IHC / OHC rows, pivoted by frequency region.
        sub = df[df["Region"].isin(["IHC region", "OHC region"])]
        if not sub.empty:
            for metric in ["Number of fibers", "Hair cells (Myo7a)",
                           "Total length (um)", "Mean diameter (um)",
                           "Branch points", "Area covered (um^2)"]:
                if metric not in sub or sub[metric].replace("", pd.NA).isna().all():
                    continue
                piv = sub.pivot_table(index="Frequency region",
                                      columns="Region", values=metric,
                                      aggfunc="mean")
                sheet = metric.split(" (")[0][:28]
                piv.to_excel(xl, sheet_name=f"{sheet}")
    return df


# --------------------------------------------------------------------------- #
#  Single-image interactive flow
# --------------------------------------------------------------------------- #

def load_and_preview(file_obj, dz, dy, dx):
    if file_obj is None:
        return (None, gr.update(), gr.update(), gr.update(),
                gr.update(), None, None, "Please upload a CZI or TIFF file.")
    try:
        img = P.load_image(file_obj, dz=float(dz), dy=float(dy), dx=float(dx))
    except Exception as e:
        return (None, gr.update(), gr.update(), gr.update(),
                gr.update(), None, None,
                f"❌ Failed to load file:\n{e}\n{traceback.format_exc()}")

    nf, myo = P.guess_channels(img)
    choices = _channel_choices(img)
    freq = P.detect_frequency(img.source_name)
    bfrac = P.suggest_boundary(img.data[myo])

    nf_prev = P.channel_preview(img.data[nf])
    myo_prev = _draw_boundary_on(P.channel_preview(img.data[myo]), bfrac, "Y")

    dz_, dy_, dx_ = img.voxel
    status = (f"✅ Loaded **{img.source_name}** — shape "
              f"{img.data.shape} (C,Z,Y,X)\n\n"
              f"Voxel size: dz={dz_:.3f}, dy={dy_:.4f}, dx={dx_:.4f} µm.  "
              f"Detected frequency region: **{freq}**.\n\n"
              f"Auto-picked Neurofilament = Ch {nf}, Myo7a = Ch {myo}. "
              f"Adjust below if needed, then press **Run analysis**.")
    return (img,
            gr.update(choices=choices, value=nf),
            gr.update(choices=choices, value=myo),
            gr.update(value=freq),
            gr.update(value=round(bfrac, 3)),
            nf_prev, myo_prev, status)


def refresh_boundary_preview(img, myo_idx, boundary, axis, hc_cents, ihc_side):
    """Redraw the Myo7a preview with the boundary and any detected hair cells."""
    if img is None or myo_idx is None:
        return None
    myo_u8 = P.channel_preview(img.data[int(myo_idx)])
    side = "low" if ihc_side.startswith("Low") else "high"
    cents = hc_cents if hc_cents is not None else np.zeros((0, 2))
    return P.hair_cell_overlay(myo_u8, cents, float(boundary), axis, side)


def detect_cells(img, myo_idx, axis):
    """Run Cellpose (or watershed fallback) on the Myo7a channel, overlay the
    detected hair cells, and propose an IHC/OHC boundary + side."""
    if img is None or myo_idx is None:
        return (None, None, gr.update(), gr.update(),
                "Load an image and select the Myo7a channel first.")
    try:
        myo_idx = int(myo_idx)
        det = P.detect_hair_cells(img.data[myo_idx], img.voxel)
        reg = P.auto_regions_from_cells(det["centroids"], det["mip_shape"])
        side_label = ("Low side = IHC" if reg["ihc_side"] == "low"
                      else "High side = IHC")
        myo_u8 = P.channel_preview(img.data[myo_idx])
        ov = P.hair_cell_overlay(myo_u8, det["centroids"],
                                 reg["boundary_frac"], axis, reg["ihc_side"])
        status = (
            f"🔬 Detected **{det['count']}** hair cells "
            f"(engine: {det['engine']}).\n\n"
            f"Suggested boundary **{reg['boundary_frac']:.3f}**, "
            f"**{side_label}**, confidence: **{reg['confidence']}** "
            f"({reg['reason']}).\n\n"
            f"⚠️ This is a *starting suggestion* — detection is often "
            f"incomplete on dense fields. **Check the overlay** (cyan = IHC, "
            f"magenta = OHC, yellow = boundary) and drag the boundary slider "
            f"to correct it before running.")
        return (det["centroids"], ov,
                gr.update(value=round(reg["boundary_frac"], 3)),
                gr.update(value=side_label), status)
    except Exception as e:
        return (None, None, gr.update(), gr.update(),
                f"❌ Hair-cell detection failed:\n{e}\n{traceback.format_exc()}")


def run_single(img, nf_idx, myo_idx, freq, axis, ihc_side, boundary,
               sensitivity, min_fiber, hc_cents):
    if img is None:
        return None, None, None, None, None, "Please load an image first."
    if nf_idx is None:
        return None, None, None, None, None, "Please select the Neurofilament channel."
    try:
        nf_idx, myo_idx = int(nf_idx), int(myo_idx)
        trace = P.trace_neurites(img.data[nf_idx], img.voxel,
                                 sensitivity=float(sensitivity))
        shape_yx = img.data[nf_idx].shape[1:]
        side = "low" if ihc_side.startswith("Low") else "high"
        ihc_roi, ohc_roi = P.make_region_masks(shape_yx, float(boundary),
                                               ihc_side=side, axis=axis)
        cents = hc_cents if (hc_cents is not None and len(hc_cents)) else None

        whole = P.compute_metrics(trace, "Whole field",
                                  min_fiber_um=float(min_fiber),
                                  hair_cell_centroids=cents)
        m_ihc = P.compute_metrics(trace, "IHC region", ihc_roi,
                                  min_fiber_um=float(min_fiber),
                                  hair_cell_centroids=cents)
        m_ohc = P.compute_metrics(trace, "OHC region", ohc_roi,
                                  min_fiber_um=float(min_fiber),
                                  hair_cell_centroids=cents)

        skel_img = P.skeleton_image(trace.skeleton)
        region_img = P.region_overlay(trace.skeleton, ihc_roi, ohc_roi,
                                      float(boundary), axis)
        myo_u8 = P.channel_preview(img.data[myo_idx])
        myo_img = P.hair_cell_overlay(
            myo_u8, cents if cents is not None else np.zeros((0, 2)),
            float(boundary), axis, side)

        stem = os.path.splitext(img.source_name)[0]
        skel_path = _save_png(skel_img, f"{stem}_skeleton.png")

        rows = []
        for m in (whole, m_ihc, m_ohc):
            r = m.as_row()
            r = {"File": img.source_name, "Frequency region": freq, **r}
            rows.append(r)
        df = pd.DataFrame(rows, columns=METRIC_COLS)
        xl_path = os.path.join(OUT_DIR, f"{stem}_quantification.xlsx")
        _build_excel(rows, xl_path)

        hc_note = ""
        if cents is not None:
            hc_note = (f" Hair cells: IHC={m_ihc.n_hair_cells}, "
                       f"OHC={m_ohc.n_hair_cells}.")
        status = (f"✅ Done. Traced {int(trace.skeleton.sum())} skeleton voxels. "
                  f"IHC={m_ihc.n_fibers} fibers / {m_ihc.total_length_um:.0f} µm, "
                  f"OHC={m_ohc.n_fibers} fibers / {m_ohc.total_length_um:.0f} µm."
                  + hc_note)
        # Return skeleton path also as downloadable file
        return (skel_img, region_img, myo_img, df,
                [skel_path, xl_path], status)
    except Exception as e:
        return (None, None, None, None, None,
                f"❌ Error:\n{e}\n{traceback.format_exc()}")


# --------------------------------------------------------------------------- #
#  Batch flow
# --------------------------------------------------------------------------- #

def run_batch(files, axis, ihc_side, sensitivity, min_fiber, detect_hc,
              dz, dy, dx, progress=gr.Progress()):
    if not files:
        return None, None, None, "Please upload one or more files."
    side = "low" if ihc_side.startswith("Low") else "high"
    all_rows, gallery, skel_paths = [], [], []
    log = []
    for f in progress.tqdm(files, desc="Processing"):
        path = f if isinstance(f, str) else f.name
        name = os.path.basename(path)
        try:
            img = P.load_image(path, dz=float(dz), dy=float(dy), dx=float(dx))
            nf, myo = P.guess_channels(img)
            freq = P.detect_frequency(name)
            trace = P.trace_neurites(img.data[nf], img.voxel,
                                     sensitivity=float(sensitivity))
            shape_yx = img.data[nf].shape[1:]
            cents = None
            hc_tag = ""
            if detect_hc:
                det = P.detect_hair_cells(img.data[myo], img.voxel)
                reg = P.auto_regions_from_cells(det["centroids"],
                                                det["mip_shape"])
                bfrac = reg["boundary_frac"]
                cur_side = reg["ihc_side"]
                cents = det["centroids"] if len(det["centroids"]) else None
                hc_tag = (f" | {det['count']} hair cells ({det['engine']}), "
                          f"auto-boundary conf={reg['confidence']}")
            else:
                bfrac = P.suggest_boundary(img.data[myo])
                cur_side = side
            ihc_roi, ohc_roi = P.make_region_masks(shape_yx, bfrac,
                                                   ihc_side=cur_side, axis=axis)
            for m in (P.compute_metrics(trace, "Whole field",
                                        min_fiber_um=float(min_fiber),
                                        hair_cell_centroids=cents),
                      P.compute_metrics(trace, "IHC region", ihc_roi,
                                        min_fiber_um=float(min_fiber),
                                        hair_cell_centroids=cents),
                      P.compute_metrics(trace, "OHC region", ohc_roi,
                                        min_fiber_um=float(min_fiber),
                                        hair_cell_centroids=cents)):
                all_rows.append({"File": name, "Frequency region": freq,
                                 **m.as_row()})
            skel_img = P.skeleton_image(trace.skeleton)
            stem = os.path.splitext(name)[0]
            sp = _save_png(skel_img, f"{stem}_skeleton.png")
            skel_paths.append(sp)
            gallery.append((skel_img, f"{name} ({freq})"))
            log.append(f"✅ {name}: {freq}{hc_tag}")
        except Exception as e:
            log.append(f"❌ {name}: {e}")

    if not all_rows:
        return None, None, gallery, "No files processed.\n" + "\n".join(log)

    xl_path = os.path.join(OUT_DIR, "batch_quantification.xlsx")
    df = _build_excel(all_rows, xl_path)

    zip_path = os.path.join(OUT_DIR, "batch_skeletons.zip")
    with zipfile.ZipFile(zip_path, "w") as z:
        for sp in skel_paths:
            z.write(sp, os.path.basename(sp))
        z.write(xl_path, os.path.basename(xl_path))

    return df, [xl_path, zip_path], gallery, "\n".join(log)


# --------------------------------------------------------------------------- #
#  UI
# --------------------------------------------------------------------------- #

INTRO = """
# 🧠 Cochlear Neurofilament Tracer

Trace auditory-nerve fibers in confocal z-stacks and quantify them **per
frequency region**, separating **IHC-innervating** from **OHC-innervating**
fibers using the Myo7a hair-cell channel.

**Channels expected:** *Neurofilament* (traces the neuron) and *Myo7a* (hair
cells — reference to split IHC vs OHC). IHCs form a single row, OHCs form three
rows, so the Myo7a band is used to place the IHC/OHC boundary — which you can
fine-tune by hand.

**Input:** Zeiss `.czi` z-stacks or generic `.tif/.tiff` stacks.
"""

with gr.Blocks(title="Cochlear Neurofilament Tracer", theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRO)

    with gr.Tab("Single image (interactive)"):
        img_state = gr.State()
        hc_state = gr.State()          # detected hair-cell centroids (Nx2)
        with gr.Row():
            with gr.Column(scale=1):
                file_in = gr.File(label="Upload CZI or TIFF",
                                  file_types=[".czi", ".tif", ".tiff"],
                                  type="filepath")
                with gr.Accordion("Voxel size (µm) — used for TIFF; CZI reads "
                                  "its own", open=False):
                    dz_in = gr.Number(0.35, label="dz (µm/plane)")
                    dy_in = gr.Number(0.0895, label="dy (µm/px)")
                    dx_in = gr.Number(0.0895, label="dx (µm/px)")
                load_btn = gr.Button("① Load & preview", variant="secondary")

                nf_dd = gr.Dropdown(label="Neurofilament channel", choices=[])
                myo_dd = gr.Dropdown(label="Myo7a channel", choices=[])
                freq_dd = gr.Dropdown(label="Frequency region",
                                      choices=P.FREQ_CHOICES,
                                      value="Other / unknown")

                gr.Markdown("**IHC / OHC region split** (Myo7a-guided)")
                axis_dd = gr.Radio(["Y", "X"], value="Y",
                                   label="Split axis (Y = radial, usual)")
                side_dd = gr.Radio(["Low side = IHC", "High side = IHC"],
                                   value="Low side = IHC",
                                   label="Which side is IHC?")
                boundary_sl = gr.Slider(0.0, 1.0, value=0.5, step=0.005,
                                        label="Boundary position (fraction "
                                              "along split axis)")
                detect_btn = gr.Button(
                    "🔬 Detect hair cells & suggest boundary "
                    f"({'Cellpose' if P.CELLPOSE_AVAILABLE else 'watershed'})",
                    variant="secondary")
                gr.Markdown(
                    "<sub>Detection is a *visual assist*: it marks hair cells "
                    "and proposes a starting boundary. Confirm/adjust against "
                    "the Myo7a overlay — it does not replace your judgement.</sub>")

                gr.Markdown("**Tracing**")
                sens_sl = gr.Slider(0.5, 1.5, value=1.0, step=0.05,
                                    label="Sensitivity (↑ = capture more/thinner "
                                          "fibers)")
                minfib_sl = gr.Slider(0.0, 20.0, value=5.0, step=0.5,
                                      label="Min fiber length to count (µm)")
                run_btn = gr.Button("② Run analysis", variant="primary")

            with gr.Column(scale=2):
                status = gr.Markdown()
                with gr.Row():
                    nf_prev = gr.Image(label="Neurofilament (MIP)",
                                       height=220)
                    myo_prev = gr.Image(label="Myo7a (MIP) + boundary",
                                        height=220)
                skel_out = gr.Image(label="Traced neurons — white on black",
                                    height=300)
                region_out = gr.Image(label="Region overlay (cyan = IHC, "
                                            "magenta = OHC)", height=300)
                table = gr.Dataframe(label="Quantification", wrap=True)
                files_out = gr.Files(label="Downloads (skeleton PNG + Excel)")

        load_btn.click(load_and_preview,
                       [file_in, dz_in, dy_in, dx_in],
                       [img_state, nf_dd, myo_dd, freq_dd, boundary_sl,
                        nf_prev, myo_prev, status]
                       ).then(lambda: None, None, hc_state)  # clear old cells
        # Live boundary preview (keeps detected hair cells visible)
        for comp in (boundary_sl, myo_dd, axis_dd, side_dd):
            comp.change(refresh_boundary_preview,
                        [img_state, myo_dd, boundary_sl, axis_dd, hc_state,
                         side_dd], myo_prev)
        detect_btn.click(detect_cells,
                         [img_state, myo_dd, axis_dd],
                         [hc_state, myo_prev, boundary_sl, side_dd, status])
        run_btn.click(run_single,
                      [img_state, nf_dd, myo_dd, freq_dd, axis_dd, side_dd,
                       boundary_sl, sens_sl, minfib_sl, hc_state],
                      [skel_out, region_out, myo_prev, table, files_out, status])

    with gr.Tab("Batch (multiple images)"):
        gr.Markdown(
            "Upload several z-stacks (e.g. all frequency regions of one "
            "cochlea). Each is auto-traced with an auto-placed IHC/OHC "
            "boundary, and results are combined into one Excel workbook "
            "organized by frequency region.")
        with gr.Row():
            with gr.Column(scale=1):
                batch_files = gr.File(label="Upload CZI/TIFF files",
                                      file_count="multiple",
                                      file_types=[".czi", ".tif", ".tiff"],
                                      type="filepath")
                b_axis = gr.Radio(["Y", "X"], value="Y", label="Split axis")
                b_side = gr.Radio(["Low side = IHC", "High side = IHC"],
                                  value="Low side = IHC",
                                  label="Which side is IHC?")
                b_sens = gr.Slider(0.5, 1.5, value=1.0, step=0.05,
                                   label="Sensitivity")
                b_minfib = gr.Slider(0.0, 20.0, value=5.0, step=0.5,
                                     label="Min fiber length (µm)")
                b_detect = gr.Checkbox(
                    value=False,
                    label="Detect hair cells & auto-place boundary "
                          f"({'Cellpose' if P.CELLPOSE_AVAILABLE else 'watershed'}"
                          ", adds ~15–20 s/image)")
                with gr.Accordion("Voxel size (µm) for TIFF", open=False):
                    b_dz = gr.Number(0.35, label="dz")
                    b_dy = gr.Number(0.0895, label="dy")
                    b_dx = gr.Number(0.0895, label="dx")
                batch_btn = gr.Button("Run batch", variant="primary")
            with gr.Column(scale=2):
                batch_log = gr.Textbox(label="Log", lines=6)
                batch_table = gr.Dataframe(label="Combined quantification",
                                           wrap=True)
                batch_files_out = gr.Files(label="Downloads (Excel + ZIP)")
                batch_gallery = gr.Gallery(label="Skeleton traces",
                                           columns=3, height=400)
        batch_btn.click(run_batch,
                        [batch_files, b_axis, b_side, b_sens, b_minfib,
                         b_detect, b_dz, b_dy, b_dx],
                        [batch_table, batch_files_out, batch_gallery, batch_log])

    gr.Markdown(
        "---\n*Method:* the Neurofilament channel is smoothed, thresholded "
        "(Otsu, scaled by the sensitivity control) and skeletonised in 3D; "
        "length, diameter (from the 3D distance transform), branch points and "
        "footprint area are measured with physical voxel spacing. Each fiber is "
        "a connected skeleton component ≥ the minimum length. The Myo7a band "
        "defines the IHC/OHC boundary, and metrics are reported for each "
        "region.\n\n"
        "*Hair-cell detection* (optional) runs **Cellpose** — or a classical "
        "watershed fallback if Cellpose isn't installed — on the Myo7a "
        "max-projection to mark hair cells, count them per region, and propose "
        "a boundary. On dense fields this detection is often incomplete, so it "
        "is offered as a **visual assist**: the numbers you trust still come "
        "from the deterministic pipeline, and the boundary remains yours to "
        "set. Detection quality improves markedly on a GPU Space.")

if __name__ == "__main__":
    demo.launch()