File size: 1,399 Bytes
753375e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
"""The ONE entrypoint a tool author edits.

`process(image_path, engine, strength)` denoises the image and returns a
noisy/denoised/clean summary + a report (PSNR/SSIM when a clean GT is available).
"""
from __future__ import annotations

import os

import numpy as np

from . import classic, viz
from .io import load_image

ENGINES = ["nlm", "tv", "wavelet", "deep"]


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

    clean = None
    side = os.path.splitext(str(image_path))[0] + "_clean.npy"
    if os.path.exists(side):
        try:
            clean = np.load(side)
        except Exception:  # noqa: BLE001
            clean = None

    if engine == "deep":
        from . import deep
        res = deep.analyze(noisy, strength=float(strength), clean=clean)
    else:
        res = classic.analyze(noisy, method=engine, strength=float(strength), clean=clean)

    return {"summary": viz.summary_image(res), "report": res["report"], "res": res}


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