core/processing.py

"""csbdeep-app — a Python runnable example of CSBDeep / CARE restoration.

CARE (Content-Aware image REstoration, Weigert et al.) restores LOW-quality
fluorescence images (low SNR + blur) toward their HIGH-quality counterpart with a
U-Net trained on paired low/high data. This app exposes it behind the standard
/process contract:

  * fast (default, no TensorFlow) — a classical CARE-style restoration: Richardson-
    Lucy deconvolution of an assumed Gaussian PSF + wavelet denoising. An honest
    stand-in for the learned network, always runnable.
  * care (wired slot) — the real csbdeep CARE model (TensorFlow), Dockerfile.care,
    weights via CARE_MODEL. Lazy-imported; not exercised in CI.

Distinct from denoise-app (denoising only) and noise2void-app (self-supervised):
CARE is *supervised content-aware restoration* (denoise + deblur).
Only this file (+ app metadata) is app-specific.
"""
from __future__ import annotations

import os
from typing import Any

import numpy as np

ENGINES = ["fast", "care"]


def _resolve(file_url: Any) -> str:
    obj = file_url
    if isinstance(obj, dict):
        obj = obj.get("path") or obj.get("url") or ""
    s = str(obj)
    if s.startswith(("http://", "https://")):
        from core.io import download_to_tmp
        return download_to_tmp(s, suffix=os.path.splitext(s)[1] or ".tif")
    return s


def _read_any(path: str) -> np.ndarray:
    ext = os.path.splitext(path)[1].lower()
    if ext in (".tif", ".tiff"):
        import tifffile
        return tifffile.imread(path)
    from PIL import Image
    return np.array(Image.open(path))


def _read_gray(file_url: Any) -> np.ndarray:
    """Read an image as grayscale float in [0, 1]."""
    arr = np.asarray(_read_any(_resolve(file_url))).astype(np.float32)
    if arr.ndim == 3:
        arr = arr[..., :3].mean(axis=-1)
    mn, mx = float(arr.min()), float(arr.max())
    return (arr - mn) / max(mx - mn, 1e-6)


def restore_fast(degraded: np.ndarray, psf_sigma: float = 2.0, strength: float = 1.0) -> np.ndarray:
    """Classical CARE-style restoration: Richardson-Lucy deconvolution + denoising."""
    from scipy.ndimage import gaussian_filter
    from skimage.restoration import denoise_wavelet, richardson_lucy

    rad = max(1, int(round(3 * psf_sigma)))
    yy, xx = np.mgrid[-rad:rad + 1, -rad:rad + 1]
    psf = np.exp(-(xx ** 2 + yy ** 2) / (2 * psf_sigma ** 2))
    psf /= psf.sum()

    pre = gaussian_filter(degraded, sigma=0.5)                   # gentle pre-smooth
    deconv = richardson_lucy(np.clip(pre, 1e-6, 1.0), psf, num_iter=15, clip=True)
    sigma = 0.08 * float(strength)
    out = denoise_wavelet(np.clip(deconv, 0, 1), sigma=sigma, rescale_sigma=True)
    return np.clip(out, 0, 1).astype(np.float32)


def _clean_sidecar(file_url: Any) -> np.ndarray | None:
    """Load the clean ground-truth sidecar (.example_clean.npy) next to a local path."""
    p = _resolve(file_url)
    if not isinstance(p, str) or p.startswith(("http://", "https://")):
        return None
    cand = os.path.join(os.path.dirname(p), ".example_clean.npy")
    if os.path.exists(cand):
        try:
            return np.load(cand).astype(np.float32)
        except Exception:  # noqa: BLE001
            return None
    return None


def _metrics(degraded, restored, clean):
    from skimage.metrics import peak_signal_noise_ratio as psnr
    from skimage.metrics import structural_similarity as ssim
    if clean is None or clean.shape != restored.shape:
        return None
    return {
        "psnr_degraded_db": round(float(psnr(clean, degraded, data_range=1.0)), 2),
        "psnr_restored_db": round(float(psnr(clean, restored, data_range=1.0)), 2),
        "ssim_degraded": round(float(ssim(clean, degraded, data_range=1.0)), 3),
        "ssim_restored": round(float(ssim(clean, restored, data_range=1.0)), 3),
    }


def process(file_url: Any, engine: str = "fast", psf_sigma: float = 2.0,
            strength: float = 1.0):
    """Restore a degraded image. Returns (degraded, restored, clean_or_None, report)."""
    if engine not in ENGINES:
        raise ValueError(f"unknown engine '{engine}'; choose {ENGINES}")
    degraded = _read_gray(file_url)
    if engine == "care":
        from core.care_engine import restore_care
        restored, model = restore_care(degraded)
    else:
        restored = restore_fast(degraded, psf_sigma=float(psf_sigma), strength=float(strength))
        model = "classical (Richardson-Lucy + wavelet)"
    clean = _clean_sidecar(file_url)
    report = {
        "engine": engine, "model": model, "shape": list(degraded.shape),
        "psf_sigma": float(psf_sigma), "strength": float(strength),
        "metrics": _metrics(degraded, restored, clean),
    }
    return degraded, restored, clean, report


def simulate_full(file_url: Any, engine: str = "fast", psf_sigma: float = 2.0,
                  strength: float = 1.0) -> dict:
    from core.viz import triptych
    degraded, restored, clean, report = process(file_url, engine=engine,
                                                psf_sigma=psf_sigma, strength=strength)
    return {"summary": triptych(degraded, restored, clean), "report": report,
            "restored": restored}