STREAMLIT_Preprocessing.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Step-1 module (callable):
- LIF loading (Leica Image File Format)
- series listing + selection
- voxel size extraction (µm/pixel)
- channel mapping (soma=blue, ion=orange)
- channel preview: single Z-slice OR MIP (Maximum Intensity Projection)
- optional preprocessing (Cellpose denoise/upsample) slice-by-slice
    -> NOW controlled by st.form + session_state cache (Submit to run)
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, Tuple, List, Any

import numpy as np
import streamlit as st


def _clear_downstream_after_preprocessing_rerun() -> None:
    # Preprocessing change invalidates segmentation + everything after it
    for k in ["cp2d_cache", "selected_slices_queue", "zones_cache_by_item", "clustering_cache", "queue_fingerprint"]:
        if k in st.session_state:
            del st.session_state[k]
    # Clear widget state for downstream steps
    for kk in list(st.session_state.keys()):
        if kk.startswith(("cp_", "zones_", "clust_", "clustering_", "queue_")):
            del st.session_state[kk]

import plotly.express as px
from readlif.reader import LifFile
from cellpose import denoise


# -------------------------
# Colorscales (blue/orange)
# -------------------------
COLOR_SOMA_BLUE = [(0.0, "black"), (1.0, "#0088FF")]   # soma = blue
COLOR_ION_ORANGE = [(0.0, "black"), (1.0, "#FF8C00")]  # ion = orange


# -------------------------
# Caching
# -------------------------
@st.cache_resource
def get_lif(path: str) -> LifFile:
    #
    # LIF files contain multiple "Series" (individual image stacks) within a single file.
    return LifFile(path)


@st.cache_resource
def get_denoise_model(model_type: str, diam_mean: float, use_gpu: bool) -> denoise.DenoiseModel:
    return denoise.DenoiseModel(gpu=use_gpu, model_type=model_type, diam_mean=float(diam_mean))


# -------------------------
# Metadata: voxel size µm/pixel
# -------------------------
def _safe_float(x):
    return float(x) if x is not None else None


def get_voxel_size_um(lif: LifFile, series_idx: int) -> Tuple[float, float, float]:
    #
    # We extract physical dimensions to ensure accurate volume calculations later.
    series = lif.get_image(series_idx)
    scale_px_per_um = series.scale  # (px/µm_x, px/µm_y, px/µm_z[, ...])

    sx = _safe_float(scale_px_per_um[0]) if len(scale_px_per_um) > 0 else None
    sy = _safe_float(scale_px_per_um[1]) if len(scale_px_per_um) > 1 else None
    sz = _safe_float(scale_px_per_um[2]) if len(scale_px_per_um) > 2 else None

    if sx is None or sy is None or sz is None or sx == 0 or sy == 0 or sz == 0:
        return (math.nan, math.nan, math.nan)

    return (1.0 / sx, 1.0 / sy, 1.0 / sz)


def list_series(lif: LifFile) -> List[Dict[str, Any]]:
    out = []
    n = len(lif.image_list)
    for i in range(n):
        s = lif.get_image(i)
        out.append({
            "idx": i,
            "name": getattr(s, "name", f"Series {i}"),
            "z": int(s.dims.z),
            "c": int(s.channels),
            "y": int(s.dims.y),
            "x": int(s.dims.x),
        })
    return out


# -------------------------
# Data loading (memory-aware)
# -------------------------
def load_slice(lif: LifFile, series_idx: int, c_idx: int, z: int) -> np.ndarray:
    s = lif.get_image(series_idx)
    arr = np.array(s.get_frame(z=int(z), t=0, c=int(c_idx)))
    return np.squeeze(arr).astype(np.float32)


def load_stack_channel(lif: LifFile, series_idx: int, c_idx: int) -> np.ndarray:
    #
    # Loads all Z-planes for a specific channel into a 3D numpy array.
    s = lif.get_image(series_idx)
    z_dim = int(s.dims.z)
    frames = [load_slice(lif, series_idx, c_idx, z) for z in range(z_dim)]
    return np.stack(frames, axis=0).astype(np.float32)


def compute_mip_channel(lif: LifFile, series_idx: int, c_idx: int) -> np.ndarray:
    #
    # Projects the 3D stack onto 2D by taking the brightest pixel through the Z-axis.
    # Useful for quick visualization of the whole cell.
    s = lif.get_image(series_idx)
    z_dim = int(s.dims.z)

    mip = None
    for z in range(z_dim):
        sl = load_slice(lif, series_idx, c_idx, z)
        mip = sl if mip is None else np.maximum(mip, sl)
    return mip.astype(np.float32)


def normalize_for_display(img2d: np.ndarray) -> np.ndarray:
    # Robust normalization using 1st and 99th percentiles to handle outliers/hot pixels.
    p1, p99 = np.percentile(img2d, 1), np.percentile(img2d, 99)
    if (p99 - p1) <= 1e-12:
        return np.clip(img2d, 0, None)
    out = (img2d - p1) / (p99 - p1)
    return np.clip(out, 0, 1)


def make_rgb_merge(
    ion2d: np.ndarray,
    soma2d: np.ndarray,
    ion_gain: float = 1.0,
    soma_gain: float = 1.0,
    gamma: float = 1.0,
) -> np.ndarray:
    #
    # Creates a composite image: Soma (Blue) + Ion (Orange/Red-Green mix).
    ion = np.clip(ion2d * float(ion_gain), 0, 1)
    soma = np.clip(soma2d * float(soma_gain), 0, 1)
    r = ion
    g = 0.55 * ion  # Mixing G into R creates Orange/Gold for Ion
    b = soma
    rgb = np.stack([r, g, b], axis=-1)
    rgb = np.clip(rgb, 0, 1)
    if gamma and gamma != 1.0:
        rgb = np.power(rgb, 1.0 / float(gamma))
    return rgb.astype(np.float32)


# -------------------------
# Preprocess
# -------------------------
@dataclass
class PreprocessChannelCfg:
    model_type: str  # "none" / "oneclick_cyto3" / "upsample_cyto3"
    diameter: float  # diam_mean


def preprocess_two_channels(
    ch_ion_zhw: np.ndarray,
    ch_soma_zhw: np.ndarray,
    ion_cfg: PreprocessChannelCfg,
    soma_cfg: PreprocessChannelCfg,
    use_gpu: bool,
) -> Tuple[np.ndarray, np.ndarray]:
    #
    # Iterates through the Z-stack and applies Cellpose denoising/upsampling model to each slice.
    z, h, w = ch_ion_zhw.shape
    out_ion = np.zeros((z, h, w), dtype=np.float32)
    out_soma = np.zeros((z, h, w), dtype=np.float32)

    ion_model = None if ion_cfg.model_type == "none" else get_denoise_model(ion_cfg.model_type, ion_cfg.diameter, use_gpu)
    soma_model = None if soma_cfg.model_type == "none" else get_denoise_model(soma_cfg.model_type, soma_cfg.diameter, use_gpu)

    for zi in range(z):
        ion_sl = ch_ion_zhw[zi]
        soma_sl = ch_soma_zhw[zi]

        if ion_model is None:
            p_ion = ion_sl
        else:
            p_ion = np.squeeze(ion_model.eval([ion_sl], channels=[0, 0], do_3D=False)[0]).astype(np.float32)

        if soma_model is None:
            p_soma = soma_sl
        else:
            p_soma = np.squeeze(soma_model.eval([soma_sl], channels=[0, 0], do_3D=False)[0]).astype(np.float32)

        if p_ion.shape != (h, w) or p_soma.shape != (h, w):
            raise ValueError(f"Preprocess shape mismatch at z={zi}: ion={p_ion.shape}, soma={p_soma.shape}, expected {(h,w)}")

        out_ion[zi] = p_ion
        out_soma[zi] = p_soma

    return out_ion, out_soma


# -------------------------
# UI helper: Step-1 panel
# -------------------------
def render_step1_lif_panel(lif_path: str) -> Dict[str, Any]:
    lif = get_lif(lif_path)
    series_list = list_series(lif)

    st.subheader("Step 1 — LIF input + channels + preview")
    st.caption(f"Series count: {len(series_list)}")

    def _fmt_series(s):
        return f"[{s['idx']}] {s['name']} | Z={s['z']} C={s['c']} {s['y']}×{s['x']}"

    series_idx = st.selectbox(
        "Select 1 series:",
        options=[s["idx"] for s in series_list],
        format_func=lambda i: _fmt_series(series_list[i]),
        key="series_idx"
    )

    s = lif.get_image(series_idx)
    n_c = int(s.channels)
    n_z = int(s.dims.z)

    vx, vy, vz = get_voxel_size_um(lif, series_idx)
    st.write(
        f"Voxel size (µm/pixel): X={vx:.4f} | Y={vy:.4f} | Z={vz:.4f}"
        if not math.isnan(vx)
        else "Voxel size (µm/pixel): **no scale metadata / nan**"
    )

    # ---- preview mode
    view_mode = st.radio("Display mode:", ["Single Z-slice", "Max intensity projection (MIP)"], horizontal=True, key="view_mode")
    z_sel = 0
    if view_mode == "Single Z-slice" and n_z > 1:
        z_sel = st.slider("Z-slice:", 0, n_z - 1, 0)

    # ---- channel roles
    st.markdown("**Channel mapping:** Soma = blue, Ion = orange. (Ignore others.)")
    roles = ["ignore", "ion (orange)", "soma (blue)"]

    role_by_c = {}
    cols = st.columns(min(n_c, 6)) if n_c > 0 else []
    for c in range(n_c):
        with cols[c % len(cols)]:
            if c == 0:
                default = 1  # ion
            elif c == 1:
                default = 2  # soma
            else:
                default = 0
            role_by_c[c] = st.selectbox(f"Ch {c}", roles, index=default, key=f"role_{series_idx}_{c}")

    ion_c = next((c for c, r in role_by_c.items() if r.startswith("ion")), None)
    soma_c = next((c for c, r in role_by_c.items() if r.startswith("soma")), None)

    if soma_c is None:
        st.error("You must select a **Soma (blue)** channel.")
        st.stop()

    if ion_c is None:
        st.warning("No Ion channel selected → pipeline will use 0-matrix later.")

    # ---- load preview images (memory-aware)
    def get_preview(c_idx: int):
        if view_mode == "Max intensity projection (MIP)":
            img = compute_mip_channel(lif, series_idx, c_idx)
        else:
            img = load_slice(lif, series_idx, c_idx, z_sel)
        return normalize_for_display(img)

    st.markdown("### Viewer (large) — Soma / Ion / Merge")

    mg1, mg2, mg3 = st.columns(3)
    with mg1:
        ion_gain = st.slider("Ion gain (merge)", 0.1, 5.0, 1.0, 0.1, key="ion_gain")
    with mg2:
        soma_gain = st.slider("Soma gain (merge)", 0.1, 5.0, 1.0, 0.1, key="soma_gain")
    with mg3:
        gamma = st.slider("Gamma (merge)", 0.5, 2.5, 1.0, 0.1, key="gamma")

    soma2d = get_preview(soma_c)
    ion2d = get_preview(ion_c) if ion_c is not None else np.zeros_like(soma2d, dtype=np.float32)
    merge_rgb = make_rgb_merge(ion2d, soma2d, ion_gain=ion_gain, soma_gain=soma_gain, gamma=gamma)

    v1, v2, v3 = st.columns(3)
    with v1:
        fig = px.imshow(soma2d, color_continuous_scale=COLOR_SOMA_BLUE,
                        title=f"SOMA (blue) — {'MIP' if view_mode.startswith('Max') else f'Z={z_sel}'}")
        fig.update_layout(height=650, margin=dict(l=0, r=0, t=40, b=0))
        fig.update_xaxes(showticklabels=False); fig.update_yaxes(showticklabels=False)
        st.plotly_chart(fig, use_container_width=True, key="plotly1")

    with v2:
        fig = px.imshow(ion2d, color_continuous_scale=COLOR_ION_ORANGE,
                        title=f"ION (orange) — {'MIP' if view_mode.startswith('Max') else f'Z={z_sel}'}")
        fig.update_layout(height=650, margin=dict(l=0, r=0, t=40, b=0))
        fig.update_xaxes(showticklabels=False); fig.update_yaxes(showticklabels=False)
        st.plotly_chart(fig, use_container_width=True, key="plotly2")

    with v3:
        fig = px.imshow(merge_rgb, title="MERGE (Soma=Blue, Ion=Orange)")
        fig.update_layout(height=650, margin=dict(l=0, r=0, t=40, b=0))
        fig.update_xaxes(showticklabels=False); fig.update_yaxes(showticklabels=False)
        st.plotly_chart(fig, use_container_width=True, key="plotly3")

    st.markdown("---")

    # =========================
    # PREPROCESS (FORM + CACHE)
    # =========================
    st.markdown("### Preprocessing (optional)")

    # A signature that depends on channel selection too
    base_sig = (str(lif_path), int(series_idx), int(soma_c), int(ion_c) if ion_c is not None else -1)

    # init cache container
    if "step1_preprocess_cache" not in st.session_state:
        st.session_state["step1_preprocess_cache"] = {}

    pp_on = st.checkbox("Preprocess ON", value=False, key=f"pp_on_{series_idx}")

    ion_proc = None
    soma_proc = None

    if not pp_on:
        st.caption("Preprocess OFF → the pipeline uses raw stacks later.")
    else:
        # 1. Initialize Session State defaults if they don't exist
        #if "pp_ion_model" not in st.session_state:
        #    st.session_state["pp_ion_model"] = "upsample_cyto3"
        #if "pp_ion_diam" not in st.session_state:
        #    st.session_state["pp_ion_diam"] = 100.0
        #if "pp_soma_model" not in st.session_state:
        #    st.session_state["pp_soma_model"] = "oneclick_cyto3"
        #if "pp_soma_diam" not in st.session_state:
        #   st.session_state["pp_soma_diam"] = 100.0

        with st.form("preprocess_form", clear_on_submit=False):
            st.write("Preprocess settings (runs only on **Submit**)")

            c1, c2 = st.columns(2)
            with c1:
                # FIXED: 'value' param removed because key="pp_use_gpu" is already in session_state, preserving user selection.
                use_gpu = st.checkbox("GPU", key="pp_use_gpu", value=True)

                # FIXED: 'index' param removed because key="pp_ion_model" handles state.
                ion_model = st.selectbox("Ion model", ["none", "upsample_cyto3", "oneclick_cyto3"],
                                         key="pp_ion_model", index=1)

                # FIXED: 'value' param removed because key="pp_ion_diam" handles state.
                ion_diam = st.number_input("Ion diameter", min_value=1.0, max_value=500.0, step=1.0,
                                           key="pp_ion_diam", value=100.0)

            with c2:
                # FIXED: 'index' param removed because key="pp_soma_model" handles state.
                soma_model = st.selectbox("Soma model", ["none", "oneclick_cyto3", "upsample_cyto3"],
                                          key="pp_soma_model", index=1)

                # FIXED: 'value' param removed because key="pp_soma_diam" handles state.
                soma_diam = st.number_input("Soma diameter", min_value=1.0, max_value=500.0, step=1.0,
                                            key="pp_soma_diam", value=100.0)

            submitted = st.form_submit_button("🚀 Run preprocessing")

        # preprocess signature (including parameters)
        pp_sig = (
            base_sig,
            bool(use_gpu),
            str(ion_model), float(ion_diam),
            str(soma_model), float(soma_diam),
        )

        cache = st.session_state["step1_preprocess_cache"]
        have_cached = pp_sig in cache

        if submitted:
            _clear_downstream_after_preprocessing_rerun()

            with st.spinner("Full Z-stack loading + preprocess..."):
                soma_stack = load_stack_channel(lif, series_idx, soma_c)
                if ion_c is None:
                    ion_stack = np.zeros_like(soma_stack, dtype=np.float32)
                else:
                    ion_stack = load_stack_channel(lif, series_idx, ion_c)

                ion_proc, soma_proc = preprocess_two_channels(
                    ch_ion_zhw=ion_stack,
                    ch_soma_zhw=soma_stack,
                    ion_cfg=PreprocessChannelCfg(model_type=str(ion_model), diameter=float(ion_diam)),
                    soma_cfg=PreprocessChannelCfg(model_type=str(soma_model), diameter=float(soma_diam)),
                    use_gpu=bool(use_gpu),
                )

                cache[pp_sig] = {
                    "ion_proc": ion_proc.astype(np.float32),
                    "soma_proc": soma_proc.astype(np.float32),
                    "params": {
                        "use_gpu": bool(use_gpu),
                        "ion_model": str(ion_model),
                        "ion_diam": float(ion_diam),
                        "soma_model": str(soma_model),
                        "soma_diam": float(soma_diam),
                    },
                }
            st.success("Preprocessing done and cached.")
            have_cached = True

        if have_cached:
            ion_proc = cache[pp_sig]["ion_proc"]
            soma_proc = cache[pp_sig]["soma_proc"]

            st.markdown("#### Preprocess viewer (large) — Soma / Ion / Merge")

            mg1, mg2, mg3 = st.columns(3)
            with mg1:
                ion_gain_pp = st.slider("Ion gain (preprocess merge)", 0.1, 5.0, 1.0, 0.1, key="ion_gain_pp")
            with mg2:
                soma_gain_pp = st.slider("Soma gain (preprocess merge)", 0.1, 5.0, 1.0, 0.1, key="soma_gain_pp")
            with mg3:
                gamma_pp = st.slider("Gamma (preprocess merge)", 0.5, 2.5, 1.0, 0.1, key="gamma_pp")

            if view_mode == "Max intensity projection (MIP)":
                ion_prev = normalize_for_display(np.max(ion_proc, axis=0))
                soma_prev = normalize_for_display(np.max(soma_proc, axis=0))
                mode_lbl = "MIP"
            else:
                ion_prev = normalize_for_display(ion_proc[z_sel])
                soma_prev = normalize_for_display(soma_proc[z_sel])
                mode_lbl = f"Z={z_sel}"

            merge_prev = make_rgb_merge(
                ion_prev, soma_prev,
                ion_gain=ion_gain_pp, soma_gain=soma_gain_pp, gamma=gamma_pp
            )

            p1, p2, p3 = st.columns(3)
            with p1:
                fig = px.imshow(ion_prev, color_continuous_scale=COLOR_ION_ORANGE,
                                title=f"ION preprocessed — {mode_lbl}")
                fig.update_layout(height=650, margin=dict(l=0, r=0, t=40, b=0))
                fig.update_xaxes(showticklabels=False); fig.update_yaxes(showticklabels=False)
                st.plotly_chart(fig, use_container_width=True, key="pp_plotly_1")

            with p2:
                fig = px.imshow(soma_prev, color_continuous_scale=COLOR_SOMA_BLUE,
                                title=f"SOMA preprocessed — {mode_lbl}")
                fig.update_layout(height=650, margin=dict(l=0, r=0, t=40, b=0))
                fig.update_xaxes(showticklabels=False); fig.update_yaxes(showticklabels=False)
                st.plotly_chart(fig, use_container_width=True, key="pp_plotly_2")

            with p3:
                fig = px.imshow(merge_prev, title=f"MERGE preprocessed — {mode_lbl}")
                fig.update_layout(height=650, margin=dict(l=0, r=0, t=40, b=0))
                fig.update_xaxes(showticklabels=False); fig.update_yaxes(showticklabels=False)
                st.plotly_chart(fig, use_container_width=True, key="pp_plotly_3")
        else:
            st.info("Press **Run preprocessing** to compute (then it will stay cached until params/input change).")

    # -------------------------
    # Output state
    # -------------------------
    state = {
        "lif_path": lif_path,
        "series_idx": int(series_idx),
        "dims": {"z": n_z, "c": n_c, "y": int(s.dims.y), "x": int(s.dims.x)},
        "channels": {"ion_c": ion_c, "soma_c": soma_c},
        "voxel_um": {"x": vx, "y": vy, "z": vz},
        "preprocess": {
            "enabled": bool(pp_on),
            "ion_stack": ion_proc if pp_on else None,
            "soma_stack": soma_proc if pp_on else None,
        },
    }
    return state