Dehaze/web_dehaze/postprocess.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image
from scipy.optimize import minimize
from skimage import color, exposure


POSTPROCESSORS: dict[str, dict[str, Any]] = {
    "manual_sv": {
        "label": "手动 S/V",
        "needs_reference": False,
        "params": {"s_gain": 1.0, "v_gain": 1.0},
    },
    "hsv_hist": {
        "label": "HSV 直方图匹配",
        "needs_reference": True,
        "params": {"match_hue": False},
    },
    "auto_sv": {
        "label": "自动 S/V",
        "needs_reference": True,
        "params": {},
    },
    "hist_auto_sv": {
        "label": "直方图 + 自动 S/V",
        "needs_reference": True,
        "params": {"match_hue": False},
    },
}


def _load_rgb_float(path: Path) -> tuple[np.ndarray, tuple[int, int]]:
    image = Image.open(path).convert("RGB")
    return np.asarray(image, dtype=np.float64) / 255.0, image.size


def _load_reference_float(path: Path, size: tuple[int, int]) -> np.ndarray:
    image = Image.open(path).convert("RGB")
    if image.size != size:
        image = image.resize(size, Image.BILINEAR)
    return np.asarray(image, dtype=np.float64) / 255.0


def _save_rgb_float(array: np.ndarray, output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image = np.clip(array, 0.0, 1.0)
    Image.fromarray((image * 255.0 + 0.5).astype(np.uint8)).save(output_path)


def adjust_sv(source_path: Path, output_path: Path, s_gain: float = 1.0, v_gain: float = 1.0) -> dict[str, Any]:
    src_rgb, _ = _load_rgb_float(source_path)
    hsv = color.rgb2hsv(src_rgb)
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * float(s_gain), 0.0, 1.0)
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] * float(v_gain), 0.0, 1.0)
    _save_rgb_float(color.hsv2rgb(hsv), output_path)
    return {"s_gain": float(s_gain), "v_gain": float(v_gain)}


def hsv_hist_match(
    source_path: Path,
    reference_path: Path,
    output_path: Path,
    match_hue: bool = False,
) -> dict[str, Any]:
    src_rgb, size = _load_rgb_float(source_path)
    ref_rgb = _load_reference_float(reference_path, size)

    hsv_src = color.rgb2hsv(src_rgb)
    hsv_ref = color.rgb2hsv(ref_rgb)

    hue = exposure.match_histograms(hsv_src[:, :, 0], hsv_ref[:, :, 0]) if match_hue else hsv_src[:, :, 0]
    sat = exposure.match_histograms(hsv_src[:, :, 1], hsv_ref[:, :, 1])
    val = exposure.match_histograms(hsv_src[:, :, 2], hsv_ref[:, :, 2])

    result_hsv = np.stack([hue, sat, val], axis=2)
    _save_rgb_float(color.hsv2rgb(result_hsv), output_path)
    return {"match_hue": bool(match_hue)}


def auto_sv(
    source_path: Path,
    reference_path: Path,
    output_path: Path,
    hist_first: bool = False,
    match_hue: bool = False,
) -> dict[str, Any]:
    src_rgb, size = _load_rgb_float(source_path)
    ref_rgb = _load_reference_float(reference_path, size)

    hsv_src = color.rgb2hsv(src_rgb)
    hsv_ref = color.rgb2hsv(ref_rgb)

    if hist_first:
        hue = exposure.match_histograms(hsv_src[:, :, 0], hsv_ref[:, :, 0]) if match_hue else hsv_src[:, :, 0]
        sat = exposure.match_histograms(hsv_src[:, :, 1], hsv_ref[:, :, 1])
        val = exposure.match_histograms(hsv_src[:, :, 2], hsv_ref[:, :, 2])
        hsv_src = np.stack([hue, sat, val], axis=2)

    def loss_function(params: np.ndarray) -> float:
        ks, kv = params
        adj_s = np.clip(hsv_src[:, :, 1] * ks, 0.0, 1.0)
        adj_v = np.clip(hsv_src[:, :, 2] * kv, 0.0, 1.0)
        loss_s = np.mean((adj_s - hsv_ref[:, :, 1]) ** 2)
        loss_v = np.mean((adj_v - hsv_ref[:, :, 2]) ** 2)
        return float(loss_s + loss_v)

    result = minimize(loss_function, np.array([1.0, 1.0]), method="Nelder-Mead", tol=1e-4)
    best_s, best_v = result.x

    hsv_final = hsv_src.copy()
    hsv_final[:, :, 1] = np.clip(hsv_final[:, :, 1] * best_s, 0.0, 1.0)
    hsv_final[:, :, 2] = np.clip(hsv_final[:, :, 2] * best_v, 0.0, 1.0)

    _save_rgb_float(color.hsv2rgb(hsv_final), output_path)
    return {
        "s_gain": float(best_s),
        "v_gain": float(best_v),
        "hist_first": bool(hist_first),
        "match_hue": bool(match_hue),
    }


def run_postprocess(
    processor: str,
    source_path: Path,
    output_path: Path,
    reference_path: Path | None = None,
    params: dict[str, Any] | None = None,
) -> dict[str, Any]:
    params = params or {}
    if processor == "manual_sv":
        return adjust_sv(
            source_path,
            output_path,
            s_gain=float(params.get("s_gain", 1.0)),
            v_gain=float(params.get("v_gain", 1.0)),
        )

    if reference_path is None:
        raise ValueError(f"{processor} requires a reference image")

    if processor == "hsv_hist":
        return hsv_hist_match(
            source_path,
            reference_path,
            output_path,
            match_hue=bool(params.get("match_hue", False)),
        )
    if processor == "auto_sv":
        return auto_sv(source_path, reference_path, output_path, hist_first=False)
    if processor == "hist_auto_sv":
        return auto_sv(
            source_path,
            reference_path,
            output_path,
            hist_first=True,
            match_hue=bool(params.get("match_hue", False)),
        )

    raise ValueError(f"Unknown postprocessor: {processor}")