core/processing.py

"""The ONE entrypoint a tool author edits.

`process(file_url, engine, radius, strength)` runs Noise2Void self-supervised
denoising on a noisy 2D image and returns a noisy/denoised/clean summary plus a
report (PSNR/SSIM when the clean ground-truth sidecar is available).
"""
from __future__ import annotations

import os

import numpy as np

from . import n2v, viz
from .io import load_image_from_any

ENGINES = n2v.ENGINES


def _to_gray01(arr: np.ndarray) -> np.ndarray:
    a = np.asarray(arr)
    if a.ndim == 3:
        a = a[..., :3].mean(-1)
    a = a.astype(np.float32)
    return np.clip((a - a.min()) / max(a.max() - a.min(), 1e-6), 0, 1)


def simulate_full(image_path, engine: str = "fast", radius: int = 2,
                  strength: float = 1.0) -> dict:
    if engine not in ENGINES:
        raise ValueError(f"Unknown engine '{engine}'. Choose one of {ENGINES}.")
    noisy = _to_gray01(load_image_from_any(image_path))

    clean = None
    # Clean ground-truth sidecar, saved next to the baked example as a hidden file
    # (leading dot) so it is never mistaken for an input image.
    if isinstance(image_path, str):
        base = os.path.splitext(image_path)[0]
        d, name = os.path.split(base)
        for side in (base + "_clean.npy", os.path.join(d, "." + name + "_clean.npy")):
            if os.path.exists(side):
                try:
                    clean = np.load(side)
                except Exception:  # noqa: BLE001
                    clean = None
                break

    res = n2v.analyze(noisy, engine=engine, radius=int(radius),
                      strength=float(strength), clean=clean)
    return {"summary": viz.summary_image(res), "report": res["report"], "res": res}


def process(image_path, engine: str = "fast", radius: int = 2,
            strength: float = 1.0) -> tuple[np.ndarray, dict]:
    """Returns (noisy/denoised/clean summary RGB, report dict)."""
    r = simulate_full(image_path, engine, radius, strength)
    return r["summary"], r["report"]