539 lines
19 KiB
Python
539 lines
19 KiB
Python
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 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}")
|
|
|
|
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)
|
|
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
|