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",
    "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", "Total length (um)",
                           "Mean diameter (um)", "Branch points",
                           "Area covered (um^2)"]:
                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):
    if img is None or myo_idx is None:
        return None
    myo_prev = P.channel_preview(img.data[int(myo_idx)])
    return _draw_boundary_on(myo_prev, float(boundary), axis)


def run_single(img, nf_idx, myo_idx, freq, axis, ihc_side, boundary,
               sensitivity, min_fiber):
    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)

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

        skel_img = P.skeleton_image(trace.skeleton)
        region_img = P.region_overlay(trace.skeleton, ihc_roi, ohc_roi,
                                      float(boundary), axis)
        myo_img = _draw_boundary_on(P.channel_preview(img.data[myo_idx]),
                                    float(boundary), axis)

        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)

        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.")
        # 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, 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))
            bfrac = P.suggest_boundary(img.data[myo])
            shape_yx = img.data[nf].shape[1:]
            ihc_roi, ohc_roi = P.make_region_masks(shape_yx, bfrac,
                                                   ihc_side=side, axis=axis)
            for m in (P.compute_metrics(trace, "Whole field",
                                        min_fiber_um=float(min_fiber)),
                      P.compute_metrics(trace, "IHC region", ihc_roi,
                                        min_fiber_um=float(min_fiber)),
                      P.compute_metrics(trace, "OHC region", ohc_roi,
                                        min_fiber_um=float(min_fiber))):
                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}")
        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()
        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)")

                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])
        # Live boundary preview
        for comp in (boundary_sl, myo_dd, axis_dd):
            comp.change(refresh_boundary_preview,
                        [img_state, myo_dd, boundary_sl, axis_dd], myo_prev)
        run_btn.click(run_single,
                      [img_state, nf_dd, myo_dd, freq_dd, axis_dd, side_dd,
                       boundary_sl, sens_sl, minfib_sl],
                      [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)")
                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_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.")

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