| """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: |
| 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"] |
|
|
