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 post_metadata_path(output_path: Path) -> Path: return output_path.with_suffix(".json") def read_post_metadata(output_path: Path) -> dict[str, Any]: metadata_path = post_metadata_path(output_path) if not metadata_path.exists(): return {} try: return json.loads(metadata_path.read_text(encoding="utf-8")) except Exception: return {} def write_post_metadata(output_path: Path, payload: dict[str, Any]) -> None: metadata_path = post_metadata_path(output_path) metadata_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") 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()): metadata = read_post_metadata(path) post.append( { "name": path.stem, "exists": True, "path": relpath(path), "meta": metadata.get("meta", {}), "processor": metadata.get("processor", ""), "source": metadata.get("source", ""), "reference": metadata.get("reference", ""), "params": metadata.get("params", {}), "metadata_path": relpath(post_metadata_path(path)) if post_metadata_path(path).exists() else "", } ) return { "image": filename, "original": {"label": "原图", "exists": True, "path": relpath(source)}, "dehaze": dehaze, "postprocess": post, } def collect_download_items(filenames: list[str] | None = None, include_original: bool = True) -> list[tuple[Path, str]]: names = filenames or [item["name"] for item in list_images()] items: list[tuple[Path, str]] = [] used_names: set[str] = set() def add_file(path: Path, archive_name: str) -> None: if not path.exists() or not path.is_file(): return archive_name = archive_name.replace("\\", "/") if archive_name in used_names: stem = Path(archive_name).with_suffix("").as_posix() suffix = Path(archive_name).suffix index = 2 while f"{stem}_{index}{suffix}" in used_names: index += 1 archive_name = f"{stem}_{index}{suffix}" used_names.add(archive_name) items.append((path, archive_name)) for filename in names: source = source_image_path(filename) folder = image_id(filename) if include_original: add_file(source, f"{folder}/original/{source.name}") for method in DEHAZE_METHODS: result = dehaze_result_path(filename, method) add_file(result, f"{folder}/dehaze/{_safe_slug(method)}.png") post_dir = image_result_dir(filename) / "postprocess" if post_dir.exists(): for result in sorted(post_dir.glob("*.png"), key=lambda p: p.name.lower()): add_file(result, f"{folder}/postprocess/{result.name}") add_file(post_metadata_path(result), f"{folder}/postprocess/{post_metadata_path(result).name}") return items 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) write_post_metadata( output, { "image": filename, "source": source_name, "processor": processor, "reference": reference_path.name, "params": proc_params, "meta": meta, "output": relpath(output), "created_at": time.strftime("%Y-%m-%d %H:%M:%S"), }, ) 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