core/deep.py

"""Deep denoiser backend — a small DnCNN trained on synthetic data at build time.

Returns the same result dict shape as `classic.analyze` so the viz + scoring are
shared. No external weights: the Dockerfile.deep trains the net at build.
"""
from __future__ import annotations

import os

import numpy as np

from . import classic, model

_NET = None


def _weights() -> str | None:
    p = os.environ.get("DENOISE_MODEL", "/app/models/dncnn.pt")
    return p if os.path.exists(p) else None


def available() -> bool:
    if not _weights():
        return False
    try:
        import torch  # noqa: F401
        return True
    except Exception:  # noqa: BLE001
        return False


def _get_net():
    global _NET
    if _NET is None:
        _NET = model.load_model(_weights())
    return _NET


def analyze(noisy: np.ndarray, strength: float = 1.0,
            clean: np.ndarray | None = None) -> dict:
    if not available():
        raise RuntimeError(
            "deep engine needs the trained DnCNN + torch. Build with Dockerfile.deep "
            "(it trains one at build). Use a classical engine (nlm/tv/wavelet) otherwise."
        )
    den = model.denoise(_get_net(), noisy)
    report = {"engine": "deep (DnCNN, trained on synthetic data)", "dims": list(noisy.shape)}
    report.update(classic.metrics(den, noisy, clean))
    return {"noisy": noisy, "denoised": den, "clean": clean, "report": report}