整合去雾网页工具

web_dehaze/pipeline.py

@@ -0,0 +1,501 @@
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any, Callable
+
+from PIL import Image
+
+from postprocess import POSTPROCESSORS, run_postprocess
+
+
+ROOT = Path(__file__).resolve().parents[1]
+IMAGE_DIR = ROOT / "待去雾图片"
+RESULTS_DIR = ROOT / "web_results"
+CONDA_ROOT = Path(os.environ.get("CONDA_EXE", sys.executable)).resolve().parents[1]
+DEFAULT_CAFFE_PYTHON = CONDA_ROOT / "envs" / "dehaze_caffe" / "bin" / "python"
+DEFAULT_TORCH_PYTHON = CONDA_ROOT / "envs" / "seg_server" / "bin" / "python"
+
+IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}
+DEHAZE_METHODS: dict[str, dict[str, Any]] = {
+    "AOD": {
+        "label": "AOD",
+        "module": "caffe",
+        "python_group": "caffe",
+        "model_file": ROOT / "AOD-Net_最好加入后处理" / "AOD_Net.caffemodel",
+    },
+    "Baidu_API": {"label": "Baidu_API", "module": "requests", "python_group": "server", "model_file": None},
+    "DCP": {"label": "DCP", "module": "cv2", "python_group": "server", "model_file": None},
+    "DehazeNet": {
+        "label": "DehazeNet",
+        "module": "caffe",
+        "python_group": "caffe",
+        "model_file": ROOT / "DehazeNet" / "DehazeNet.caffemodel",
+    },
+    "GCANet": {
+        "label": "GCANet",
+        "module": "torch",
+        "python_group": "torch",
+        "model_file": ROOT / "GCANet" / "models" / "wacv_gcanet_dehaze.pth",
+    },
+    "RefineDNet": {
+        "label": "RefineDNet",
+        "module": "torch",
+        "python_group": "torch",
+        "model_file": ROOT / "RefineDNet" / "checkpoints" / "refined_DCP_outdoor" / "60_net_Refiner_J.pth",
+    },
+}
+
+LogFn = Callable[[str], None]
+
+
+def _noop_log(_: str) -> None:
+    return None
+
+
+def _safe_slug(value: str) -> str:
+    slug = re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._-")
+    return slug or "image"
+
+
+def image_id(filename: str) -> str:
+    stem = Path(filename).stem
+    digest = hashlib.sha1(filename.encode("utf-8")).hexdigest()[:8]
+    return f"{_safe_slug(stem)}_{digest}"
+
+
+def source_image_path(filename: str) -> Path:
+    if Path(filename).name != filename:
+        raise ValueError("Invalid image filename")
+    path = IMAGE_DIR / filename
+    if not path.exists() or path.suffix.lower() not in IMAGE_EXTENSIONS:
+        raise FileNotFoundError(filename)
+    return path
+
+
+def image_result_dir(filename: str) -> Path:
+    return RESULTS_DIR / image_id(filename)
+
+
+def dehaze_result_path(filename: str, method: str) -> Path:
+    return image_result_dir(filename) / "dehaze" / f"{_safe_slug(method)}.png"
+
+
+def post_result_path(filename: str, source: str, processor: str, params: dict[str, Any] | None = None) -> Path:
+    params = params or {}
+    suffix = ""
+    if processor == "manual_sv":
+        suffix = f"_S{int(float(params.get('s_gain', 1.0)) * 100)}_V{int(float(params.get('v_gain', 1.0)) * 100)}"
+    if params.get("match_hue"):
+        suffix += "_H"
+    return image_result_dir(filename) / "postprocess" / f"{_safe_slug(source)}__{_safe_slug(processor)}{suffix}.png"
+
+
+def relpath(path: Path) -> str:
+    return path.resolve().relative_to(ROOT).as_posix()
+
+
+def python_for_group(group: str | None) -> Path:
+    if group == "caffe":
+        return Path(os.environ.get("DEHAZE_CAFFE_PYTHON", DEFAULT_CAFFE_PYTHON))
+    if group == "torch":
+        return Path(os.environ.get("DEHAZE_TORCH_PYTHON", DEFAULT_TORCH_PYTHON))
+    return Path(sys.executable)
+
+
+def python_for_method(method: str) -> Path:
+    return python_for_group(DEHAZE_METHODS[method].get("python_group"))
+
+
+def list_images() -> list[dict[str, Any]]:
+    images: list[dict[str, Any]] = []
+    if not IMAGE_DIR.exists():
+        return images
+    for path in sorted(IMAGE_DIR.iterdir(), key=lambda p: p.name.lower()):
+        if not path.is_file() or path.suffix.lower() not in IMAGE_EXTENSIONS:
+            continue
+        width = height = None
+        mode = ""
+        try:
+            with Image.open(path) as img:
+                width, height = img.size
+                mode = img.mode
+        except Exception:
+            pass
+        images.append(
+            {
+                "name": path.name,
+                "id": image_id(path.name),
+                "size": path.stat().st_size,
+                "width": width,
+                "height": height,
+                "mode": mode,
+                "path": relpath(path),
+            }
+        )
+    return images
+
+
+def module_available(module_name: str, python_path: Path | None = None) -> bool:
+    python_path = python_path or Path(sys.executable)
+    if python_path.resolve() == Path(sys.executable).resolve():
+        return importlib.util.find_spec(module_name) is not None
+    if not python_path.exists():
+        return False
+    result = subprocess.run(
+        [str(python_path), "-c", f"import {module_name}"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        text=True,
+        timeout=20,
+    )
+    return result.returncode == 0
+
+
+def capabilities() -> dict[str, Any]:
+    methods = {}
+    for key, info in DEHAZE_METHODS.items():
+        module_name = info.get("module")
+        model_file = info.get("model_file")
+        python_path = python_for_method(key)
+        python_ok = python_path.exists()
+        module_ok = bool(module_available(module_name, python_path)) if module_name and python_ok else False
+        model_ok = bool(model_file.exists()) if model_file else True
+        methods[key] = {
+            "label": info["label"],
+            "available": python_ok and module_ok and model_ok,
+            "module": module_name,
+            "python": str(python_path),
+            "python_ok": python_ok,
+            "module_ok": module_ok,
+            "model_ok": model_ok,
+            "model_file": relpath(model_file) if model_file else "",
+        }
+    return {
+        "python": sys.executable,
+        "image_dir": relpath(IMAGE_DIR),
+        "results_dir": relpath(RESULTS_DIR),
+        "methods": methods,
+        "postprocessors": POSTPROCESSORS,
+    }
+
+
+def get_results(filename: str) -> dict[str, Any]:
+    source = source_image_path(filename)
+    dehaze = []
+    for method, info in DEHAZE_METHODS.items():
+        path = dehaze_result_path(filename, method)
+        dehaze.append(
+            {
+                "method": method,
+                "label": info["label"],
+                "exists": path.exists(),
+                "path": relpath(path) if path.exists() else "",
+            }
+        )
+
+    post = []
+    post_dir = image_result_dir(filename) / "postprocess"
+    if post_dir.exists():
+        for path in sorted(post_dir.glob("*.png"), key=lambda p: p.name.lower()):
+            post.append({"name": path.stem, "exists": True, "path": relpath(path)})
+
+    return {
+        "image": filename,
+        "original": {"label": "原图", "exists": True, "path": relpath(source)},
+        "dehaze": dehaze,
+        "postprocess": post,
+    }
+
+
+def _reset_dir(path: Path) -> None:
+    if path.exists():
+        shutil.rmtree(path)
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def _prepare_rgb_png(source: Path, target_dir: Path, min_side: int | None = None) -> Path:
+    target_dir.mkdir(parents=True, exist_ok=True)
+    image = Image.open(source).convert("RGB")
+    if min_side and min(image.size) < min_side:
+        scale = min_side / float(min(image.size))
+        new_size = (max(1, int(round(image.size[0] * scale))), max(1, int(round(image.size[1] * scale))))
+        image = image.resize(new_size, Image.BICUBIC)
+    target = target_dir / f"{_safe_slug(source.stem)}.png"
+    image.save(target)
+    return target
+
+
+def _copy_final_image(source: Path, destination: Path, original: Path, force_original_size: bool = True) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    if not force_original_size:
+        shutil.copy2(source, destination)
+        return
+    original_size = Image.open(original).size
+    image = Image.open(source).convert("RGB")
+    if image.size != original_size:
+        image = image.resize(original_size, Image.BICUBIC)
+    image.save(destination)
+
+
+def _run_command(command: list[str], cwd: Path, log: LogFn, timeout: int | None = None) -> None:
+    log(f"$ {' '.join(command)}")
+    env = os.environ.copy()
+    env.setdefault("GLOG_minloglevel", "2")
+    process = subprocess.Popen(
+        command,
+        cwd=str(cwd),
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
+    start = time.time()
+    assert process.stdout is not None
+    for line in process.stdout:
+        log(line.rstrip())
+        if timeout and time.time() - start > timeout:
+            process.kill()
+            raise TimeoutError(f"Command timed out after {timeout}s")
+    return_code = process.wait()
+    if return_code != 0:
+        raise RuntimeError(f"Command failed with code {return_code}")
+
+
+def _require_module(module_name: str, python_path: Path | None = None) -> None:
+    python_path = python_path or Path(sys.executable)
+    if not module_available(module_name, python_path):
+        raise RuntimeError(f"Python 环境 {python_path} 缺少模块：{module_name}")
+
+
+def run_dehaze_method(filename: str, method: str, options: dict[str, Any] | None = None, log: LogFn | None = None) -> Path:
+    options = options or {}
+    log = log or _noop_log
+    if method not in DEHAZE_METHODS:
+        raise ValueError(f"Unknown method: {method}")
+
+    source = source_image_path(filename)
+    log(f"开始 {method}: {filename}")
+
+    runners = {
+        "DCP": _run_dcp,
+        "Baidu_API": _run_baidu,
+        "AOD": _run_aod,
+        "DehazeNet": _run_dehazenet,
+        "GCANet": _run_gcanet,
+        "RefineDNet": _run_refinednet,
+    }
+    result = runners[method](source, filename, options, log)
+    log(f"完成 {method}: {relpath(result)}")
+    return result
+
+
+def _run_dcp(source: Path, filename: str, options: dict[str, Any], log: LogFn) -> Path:
+    _require_module("cv2")
+    work = image_result_dir(filename) / "work" / "DCP"
+    _reset_dir(work)
+    src_dir = work / "src"
+    (work / "dark").mkdir(parents=True, exist_ok=True)
+    (work / "trans").mkdir(parents=True, exist_ok=True)
+    (work / "result").mkdir(parents=True, exist_ok=True)
+    input_copy = _prepare_rgb_png(source, src_dir)
+    sz = int(options.get("sz", 10))
+    tx = float(options.get("tx", 0.2))
+    _run_command([sys.executable, "dehaze.py", str(work), str(sz), str(tx)], ROOT / "DCP_最好加入后处理", log)
+    generated = work / "result" / f"{input_copy.stem}_{sz}_{tx}_result.png"
+    if not generated.exists():
+        matches = sorted((work / "result").glob("*_result.png"))
+        if not matches:
+            raise FileNotFoundError("DCP did not create a result image")
+        generated = matches[-1]
+    output = dehaze_result_path(filename, "DCP")
+    _copy_final_image(generated, output, source)
+    return output
+
+
+def _run_baidu(source: Path, filename: str, options: dict[str, Any], log: LogFn) -> Path:
+    _require_module("requests")
+    script_path = ROOT / "Baidu_API_最好加入后处理" / "1_Baidu_Dehaze.py"
+    spec = importlib.util.spec_from_file_location("baidu_dehaze_script", script_path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError("无法加载 Baidu API 脚本")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    if os.environ.get("BAIDU_API_KEY"):
+        module.API_KEY = os.environ["BAIDU_API_KEY"]
+    if os.environ.get("BAIDU_SECRET_KEY"):
+        module.SECRET_KEY = os.environ["BAIDU_SECRET_KEY"]
+
+    token = module.get_access_token()
+    if not token:
+        raise RuntimeError("Baidu access token 获取失败")
+    log("Baidu access token 获取成功")
+    processed = module.process_image(str(source), token)
+    if not processed:
+        raise RuntimeError("Baidu API 未返回图像")
+
+    work = image_result_dir(filename) / "work" / "Baidu_API"
+    _reset_dir(work)
+    tmp = work / "baidu_result.png"
+    tmp.write_bytes(processed)
+    output = dehaze_result_path(filename, "Baidu_API")
+    _copy_final_image(tmp, output, source)
+    return output
+
+
+def _run_aod(source: Path, filename: str, options: dict[str, Any], log: LogFn) -> Path:
+    python_path = python_for_method("AOD")
+    _require_module("caffe", python_path)
+    work = image_result_dir(filename) / "work" / "AOD"
+    _reset_dir(work)
+    input_dir = work / "input"
+    output_dir = work / "output"
+    input_copy = _prepare_rgb_png(source, input_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    _run_command([str(python_path), "test/test.py", str(input_dir), str(output_dir)], ROOT / "AOD-Net_最好加入后处理", log)
+    generated = output_dir / f"{input_copy.stem}_AOD-Net.png"
+    if not generated.exists():
+        matches = sorted(output_dir.glob("*.png"))
+        if not matches:
+            raise FileNotFoundError("AOD did not create a result image")
+        generated = matches[-1]
+    output = dehaze_result_path(filename, "AOD")
+    _copy_final_image(generated, output, source)
+    return output
+
+
+def _run_dehazenet(source: Path, filename: str, options: dict[str, Any], log: LogFn) -> Path:
+    python_path = python_for_method("DehazeNet")
+    _require_module("caffe", python_path)
+    work = image_result_dir(filename) / "work" / "DehazeNet" / "img"
+    _reset_dir(work)
+    input_copy = _prepare_rgb_png(source, work / "src")
+    _run_command([str(python_path), "DehazeNet.py", str(work)], ROOT / "DehazeNet", log)
+    generated = work / "result" / f"{input_copy.stem}_result.png"
+    if not generated.exists():
+        matches = sorted((work / "result").glob("*_result.png"))
+        if not matches:
+            raise FileNotFoundError("DehazeNet did not create a result image")
+        generated = matches[-1]
+    output = dehaze_result_path(filename, "DehazeNet")
+    _copy_final_image(generated, output, source)
+    return output
+
+
+def _run_gcanet(source: Path, filename: str, options: dict[str, Any], log: LogFn) -> Path:
+    python_path = python_for_method("GCANet")
+    _require_module("torch", python_path)
+    work = image_result_dir(filename) / "work" / "GCANet"
+    _reset_dir(work)
+    input_dir = work / "input"
+    output_dir = work / "output"
+    input_copy = _prepare_rgb_png(source, input_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    _run_command(
+        [str(python_path), "test.py", "--task", "dehaze", "--gpu_id", "0", "--indir", str(input_dir), "--outdir", str(output_dir)],
+        ROOT / "GCANet",
+        log,
+    )
+    generated = output_dir / f"{input_copy.stem}_dehaze.png"
+    if not generated.exists():
+        matches = sorted(output_dir.glob("*.png"))
+        if not matches:
+            raise FileNotFoundError("GCANet did not create a result image")
+        generated = matches[-1]
+    output = dehaze_result_path(filename, "GCANet")
+    _copy_final_image(generated, output, source)
+    return output
+
+
+def _run_refinednet(source: Path, filename: str, options: dict[str, Any], log: LogFn) -> Path:
+    python_path = python_for_method("RefineDNet")
+    _require_module("torch", python_path)
+    _require_module("torchvision", python_path)
+    work = image_result_dir(filename) / "work" / "RefineDNet"
+    _reset_dir(work)
+    dataroot = work / "dataset"
+    input_copy = _prepare_rgb_png(source, dataroot / "test", min_side=256)
+    method_name = "RefineDNet"
+    _run_command(
+        [
+            str(python_path),
+            "quick_test.py",
+            "--dataroot",
+            str(dataroot),
+            "--dataset_mode",
+            "single",
+            "--name",
+            "refined_DCP_outdoor",
+            "--model",
+            "refined_DCP",
+            "--phase",
+            "test",
+            "--preprocess",
+            "none",
+            "--save_image",
+            "--method_name",
+            method_name,
+            "--epoch",
+            "60",
+            "--gpu_ids",
+            "0",
+        ],
+        ROOT / "RefineDNet",
+        log,
+    )
+    generated = dataroot / method_name / f"{input_copy.stem}_dehz.png"
+    if not generated.exists():
+        matches = sorted((dataroot / method_name).glob("*.png"))
+        if not matches:
+            raise FileNotFoundError("RefineDNet did not create a result image")
+        generated = matches[-1]
+    output = dehaze_result_path(filename, "RefineDNet")
+    _copy_final_image(generated, output, source)
+    return output
+
+
+def run_postprocessors_for_source(
+    filename: str,
+    source_name: str,
+    processors: list[str],
+    params: dict[str, Any] | None = None,
+    reference_filename: str | None = None,
+    log: LogFn | None = None,
+) -> list[Path]:
+    params = params or {}
+    log = log or _noop_log
+    if source_name == "original":
+        source_path = source_image_path(filename)
+    else:
+        source_path = dehaze_result_path(filename, source_name)
+    if not source_path.exists():
+        raise FileNotFoundError(f"后处理源图不存在：{source_name}")
+
+    reference_path = source_image_path(reference_filename or filename)
+    outputs: list[Path] = []
+    for processor in processors:
+        proc_params = params.get(processor, params)
+        output = post_result_path(filename, source_name, processor, proc_params)
+        meta = run_postprocess(processor, source_path, output, reference_path=reference_path, params=proc_params)
+        outputs.append(output)
+        log(f"后处理完成 {source_name} / {processor}: {json.dumps(meta, ensure_ascii=False)}")
+    return outputs
+
+
+def resolve_asset(relative_path: str) -> Path:
+    target = (ROOT / relative_path).resolve()
+    allowed_roots = [IMAGE_DIR.resolve(), RESULTS_DIR.resolve()]
+    if not any(target == root or root in target.parents for root in allowed_roots):
+        raise PermissionError("Asset path is outside allowed roots")
+    if not target.exists() or not target.is_file():
+        raise FileNotFoundError(relative_path)
+    return target