core/processing.py

"""The ONE entrypoint a tool author edits.

DeepInterpolation-style self-supervised denoising of a fluorescence movie.
`process(movie, ...)` predicts each frame from its temporal neighbors (center
excluded), returns a noisy/denoised/clean mid-frame triptych (RGB) plus a report
with PSNR/SSIM when the clean ground truth sidecar is present.

Two engines behind one contract:
  * fast              — classical Gaussian-weighted temporal-neighbor predictor
                        (default, always runnable; DeepInterpolation minus the net)
  * deepinterpolation — real Allen TensorFlow package (Dockerfile.deepinterp)
"""
from __future__ import annotations

import os

import numpy as np

from . import deepinterp_engine, metrics, neighbor, viz
from .io import APP_TMP_DIR

ENGINES = ["fast", "deepinterpolation"]


def _to_movie(arr: np.ndarray) -> np.ndarray:
    arr = np.asarray(arr, dtype=np.float32)
    if arr.ndim == 2:
        arr = arr[None]
    if arr.ndim != 3:
        raise ValueError(f"Expected a (T, H, W) movie, got shape {arr.shape}.")
    return arr


def _find_clean_sidecar(source_path: str | None) -> np.ndarray | None:
    """Load the clean ground-truth movie saved next to the example, if any."""
    candidates = []
    if source_path:
        base = os.path.splitext(str(source_path))[0]
        candidates.append(base + "_clean.npy")
    candidates.append(str(APP_TMP_DIR / "example_clean.npy"))
    for c in candidates:
        if c and os.path.exists(c):
            try:
                return np.asarray(np.load(c), dtype=np.float32)
            except Exception:  # noqa: BLE001
                continue
    return None


def simulate_full(movie: np.ndarray, engine: str = "fast", pre: int = 5,
                  post: int = 5, omit: int = 0,
                  source_path: str | None = None) -> dict:
    if engine not in ENGINES:
        raise ValueError(f"Unknown engine '{engine}'. Choose one of {ENGINES}.")
    noisy = _to_movie(movie)

    if engine == "deepinterpolation":
        denoised = deepinterp_engine.interpolate(noisy, pre=int(pre), post=int(post),
                                                 omit=int(omit))
        eng = "deepinterpolation (Allen, learned)"
    else:
        denoised = neighbor.interpolate(noisy, pre=int(pre), post=int(post),
                                        omit=int(omit))
        eng = "fast (Gaussian temporal-neighbor predictor)"

    T, H, W = noisy.shape
    report = {
        "engine": eng,
        "n_frames": int(T),
        "dims": [int(H), int(W)],
        "pre": int(pre),
        "post": int(post),
        "omit": int(omit),
    }

    clean = _find_clean_sidecar(source_path)
    if clean is not None and clean.shape == noisy.shape:
        report["psnr_noisy_db"] = metrics.movie_psnr(clean, noisy)
        report["psnr_denoised_db"] = metrics.movie_psnr(clean, denoised)
        report["ssim_noisy"] = metrics.movie_ssim(clean, noisy)
        report["ssim_denoised"] = metrics.movie_ssim(clean, denoised)
        report["psnr_gain_db"] = report["psnr_denoised_db"] - report["psnr_noisy_db"]
    else:
        clean = None

    frame = int(T // 2)
    summary = viz.triptych(noisy, denoised, clean, frame, report)

    return {
        "summary": summary,
        "report": report,
        "noisy": noisy,
        "denoised": denoised,
        "clean": clean,
        "frame": frame,
    }


def process(movie: np.ndarray, engine: str = "fast", pre: int = 5, post: int = 5,
            omit: int = 0, source_path: str | None = None) -> tuple[np.ndarray, dict]:
    """Returns (noisy/denoised/clean triptych RGB, report dict)."""
    r = simulate_full(movie, engine=engine, pre=pre, post=post, omit=omit,
                      source_path=source_path)
    return r["summary"], r["report"]