整合去雾网页工具
This commit is contained in:
admin
79
AOD-Net_最好加入后处理/test/test.py
Normal file
79
AOD-Net_最好加入后处理/test/test.py
Normal file
@@ -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()
|
||||
189
AOD-Net_最好加入后处理/test/test_template.prototxt
Normal file
189
AOD-Net_最好加入后处理/test/test_template.prototxt
Normal file
@@ -0,0 +1,189 @@
|
||||
name:"DirectDehazing"
|
||||
input: "data"
|
||||
input_dim: 1
|
||||
input_dim: 3
|
||||
input_dim:{height}
|
||||
input_dim:{width}
|
||||
|
||||
input: "label"
|
||||
input_dim: 1
|
||||
input_dim: 3
|
||||
input_dim:{height}
|
||||
input_dim:{width}
|
||||
|
||||
|
||||
layer {{
|
||||
name: "conv1"
|
||||
type: "Convolution"
|
||||
bottom: "data"
|
||||
top: "conv1"
|
||||
convolution_param {{
|
||||
num_output: 3
|
||||
kernel_size: 1
|
||||
}}
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "relu1"
|
||||
type: "ReLU"
|
||||
bottom: "conv1"
|
||||
top: "conv1"
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "conv2"
|
||||
type: "Convolution"
|
||||
bottom: "conv1"
|
||||
top: "conv2"
|
||||
convolution_param {{
|
||||
num_output: 3
|
||||
kernel_size: 3
|
||||
pad:1
|
||||
}}
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "relu2"
|
||||
type: "ReLU"
|
||||
bottom: "conv2"
|
||||
top: "conv2"
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "Concat1"
|
||||
type: "Concat"
|
||||
bottom: "conv1"
|
||||
bottom: "conv2"
|
||||
top: "Concat1"
|
||||
concat_param {{
|
||||
axis: 1
|
||||
}}
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "conv3"
|
||||
type: "Convolution"
|
||||
bottom: "Concat1"
|
||||
top: "conv3"
|
||||
convolution_param {{
|
||||
num_output: 3
|
||||
kernel_size: 5
|
||||
pad:2
|
||||
}}
|
||||
}}
|
||||
layer {{
|
||||
name: "relu3"
|
||||
type: "ReLU"
|
||||
bottom: "conv3"
|
||||
top: "conv3"
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "Concat2"
|
||||
type: "Concat"
|
||||
bottom: "conv2"
|
||||
bottom: "conv3"
|
||||
top: "Concat2"
|
||||
concat_param {{
|
||||
axis: 1
|
||||
}}
|
||||
}}
|
||||
layer {{
|
||||
name: "conv4"
|
||||
type: "Convolution"
|
||||
bottom: "Concat2"
|
||||
top: "conv4"
|
||||
convolution_param {{
|
||||
num_output: 3
|
||||
kernel_size: 7
|
||||
pad:3
|
||||
}}
|
||||
}}
|
||||
layer {{
|
||||
name: "relu4"
|
||||
type: "ReLU"
|
||||
bottom: "conv4"
|
||||
top: "conv4"
|
||||
}}
|
||||
layer {{
|
||||
name: "Concat3"
|
||||
type: "Concat"
|
||||
bottom: "conv1"
|
||||
bottom: "conv2"
|
||||
bottom: "conv3"
|
||||
bottom: "conv4"
|
||||
top: "Concat3"
|
||||
concat_param {{
|
||||
axis: 1
|
||||
}}
|
||||
}}
|
||||
layer {{
|
||||
name: "conv5"
|
||||
type: "Convolution"
|
||||
bottom: "Concat3"
|
||||
top: "conv5"
|
||||
convolution_param {{
|
||||
num_output: 3
|
||||
kernel_size: 3
|
||||
pad:1
|
||||
}}
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "relu5"
|
||||
type: "ReLU"
|
||||
bottom: "conv5"
|
||||
top: "K"
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name: "prod"
|
||||
type: "Eltwise"
|
||||
bottom: "data"
|
||||
bottom: "K"
|
||||
top: "prod"
|
||||
eltwise_param {{
|
||||
operation: PROD
|
||||
}}
|
||||
}}
|
||||
|
||||
layer {{
|
||||
name:"eltwise_layer"
|
||||
type:"Eltwise"
|
||||
bottom:"prod"
|
||||
bottom:"K"
|
||||
top:"eltwise_layer"
|
||||
eltwise_param{{
|
||||
operation:SUM
|
||||
coeff:1
|
||||
coeff:-1
|
||||
}}
|
||||
}}
|
||||
|
||||
|
||||
layer {{
|
||||
name: "sum"
|
||||
bottom: "eltwise_layer"
|
||||
top: "sum"
|
||||
type: "Power"
|
||||
power_param {{
|
||||
power: 1
|
||||
scale: 1
|
||||
shift: 1
|
||||
}}
|
||||
}}
|
||||
layer {{
|
||||
name: "clip"
|
||||
type: "ReLU"
|
||||
bottom: "sum"
|
||||
top: "sum"
|
||||
}}
|
||||
layer {{
|
||||
name: "loss"
|
||||
type: "EuclideanLoss"
|
||||
bottom: "sum"
|
||||
bottom: "label"
|
||||
top: "loss"
|
||||
}}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user