整合去雾网页工具

AOD-Net_最好加入后处理/test/test.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import caffe
+import cv2
+
+
+BASE_DIR = Path(__file__).resolve().parents[1]
+SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
+
+
+def edit_fcn_proto(template_file: Path, output_file: Path, height: int, width: int) -> None:
+    template = template_file.read_text()
+    output_file.write_text(template.format(height=height, width=width))
+
+
+def iter_images(input_dir: Path) -> list[Path]:
+    return [
+        path
+        for path in sorted(input_dir.iterdir(), key=lambda p: p.name.lower())
+        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS
+    ]
+
+
+def run(input_dir: Path, output_dir: Path, max_dim: int) -> None:
+    caffe.set_mode_cpu()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    template_file = BASE_DIR / "test" / "test_template.prototxt"
+    deploy_file = BASE_DIR / "test" / "DeployT.prototxt"
+    model_file = BASE_DIR / "AOD_Net.caffemodel"
+
+    images = iter_images(input_dir)
+    print(f"image numbers: {len(images)}")
+
+    for index, img_path in enumerate(images, start=1):
+        npstore = caffe.io.load_image(str(img_path))
+        orig_h, orig_w = npstore.shape[0], npstore.shape[1]
+
+        if orig_h > max_dim or orig_w > max_dim:
+            scale = max_dim / float(max(orig_h, orig_w))
+            new_h = int(orig_h * scale)
+            new_w = int(orig_w * scale)
+            npstore = cv2.resize(npstore, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
+            print(f"Resized {img_path.name} from {orig_w}x{orig_h} to {new_w}x{new_h}")
+
+        height, width = npstore.shape[0], npstore.shape[1]
+        edit_fcn_proto(template_file, deploy_file, height, width)
+
+        net = caffe.Net(str(deploy_file), str(model_file), caffe.TEST)
+        data = npstore.transpose((2, 0, 1))
+        net.blobs["data"].data[...] = [data]
+        net.forward()
+
+        result = net.blobs["sum"].data[0].transpose((1, 2, 0))
+        result = result[:, :, ::-1]
+
+        if height != orig_h or width != orig_w:
+            result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_CUBIC)
+
+        save_path = output_dir / f"{img_path.stem}_AOD-Net.png"
+        cv2.imwrite(str(save_path), result * 255.0, [cv2.IMWRITE_JPEG_QUALITY, 100])
+        print(f"[{index}/{len(images)}] saved: {save_path}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run AOD-Net on a folder of images.")
+    parser.add_argument("input_dir", nargs="?", default=str(BASE_DIR / "data" / "img"))
+    parser.add_argument("output_dir", nargs="?", default=str(BASE_DIR / "data" / "result"))
+    parser.add_argument("--max-dim", type=int, default=1920)
+    args = parser.parse_args()
+
+    run(Path(args.input_dir), Path(args.output_dir), args.max_dim)
+
+
+if __name__ == "__main__":
+    main()