整合去雾网页工具

web_dehaze/postprocess.py

@@ -0,0 +1,164 @@
+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}")