整合去雾网页工具

This commit is contained in:
admin
2026-06-10 17:42:11 +08:00
commit 6db15ebc3f
101 changed files with 10167 additions and 0 deletions

30
.gitignore vendored Normal file
View File

@@ -0,0 +1,30 @@
# Python
__pycache__/
*.py[cod]
*.pyo
# OS/editor noise
.DS_Store
Thumbs.db
.vscode/
.idea/
# Local credentials and runtime config
.env
config/*.env
# Generated dehaze outputs and temporary model inputs
web_results/
AOD-Net_最好加入后处理/data/img/
AOD-Net_最好加入后处理/data/result/
DCP_最好加入后处理/image/
DehazeNet/img/
GCANet/imgs/
RefineDNet/datasets/quick_test/
# Runtime-generated prototxt files
AOD-Net_最好加入后处理/test/DeployT.prototxt
DehazeNet/DehazeNetFcn.prototxt
# Logs
*.log

View File

@@ -0,0 +1,24 @@
We appreciate the [PyTorch implementation](https://github.com/TheFairBear/PyTorch-Image-Dehazing) of AOD-Net. Based on this code, we add simple [PONO and MS](https://github.com/Boyiliee/PONO) into AOD-Net, which improves the performance efficiently.
Previous AOD-Net Results:
![](../AOD-Net_result.png)
For TestSet A, the PSNR increases from 19.69 to 20.38 dB, the SSIM increases from 0.8478 to 0.8587. For TestSetB, the PSNR increases from 21.54 to 21.67 dB, the SSIM increases from 0.9272 to 0.9285.
If you find this repo useful, please cite:
```
@inproceedings{ICCV17a,
title={AOD-Net: All-in-One Dehazing Network},
author={Li, Boyi and Peng, Xiulian and Wang, Zhangyang and Xu, Ji-Zheng and Feng, Dan},
booktitle={Proceedings of the IEEE International Conference on Computer Vision},
year={2017}
}
@inproceedings{li2019positional,
title={Positional Normalization},
author={Li, Boyi and Wu, Felix and Weinberger, Kilian Q and Belongie, Serge},
booktitle={Advances in Neural Information Processing Systems},
pages={1620--1632},
year={2019}
}
```

View File

@@ -0,0 +1,99 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class AODnet(nn.Module):
def __init__(self):
super(AODnet, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=1, stride=1, padding=0)
self.conv2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1)
self.conv3 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=5, stride=1, padding=2)
self.conv4 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=7, stride=1, padding=3)
self.conv5 = nn.Conv2d(in_channels=12, out_channels=3, kernel_size=3, stride=1, padding=1)
self.b = 1
def forward(self, x):
x1 = F.relu(self.conv1(x))
x2 = F.relu(self.conv2(x1))
cat1 = torch.cat((x1, x2), 1)
x3 = F.relu(self.conv3(cat1))
cat2 = torch.cat((x2, x3), 1)
x4 = F.relu(self.conv4(cat2))
cat3 = torch.cat((x1, x2, x3, x4), 1)
k = F.relu(self.conv5(cat3))
if k.size() != x.size():
raise Exception("k, haze image are different size!")
output = k * x - k + self.b
return F.relu(output)
class AOD_pono_net(nn.Module):
def __init__(self):
super(AOD_pono_net, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=1, stride=1, padding=0)
self.conv2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1)
self.conv3 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=5, stride=1, padding=2)
self.conv4 = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=7, stride=1, padding=3)
self.conv5 = nn.Conv2d(in_channels=12, out_channels=3, kernel_size=3, stride=1, padding=1)
self.b = 1
self.pono = PONO(affine=False)
self.ms = MS()
def forward(self, x):
x1 = F.relu(self.conv1(x))
x2 = F.relu(self.conv2(x1))
cat1 = torch.cat((x1, x2), 1)
x1, mean1, std1 = self.pono(x1)
x2, mean2, std2 = self.pono(x2)
x3 = F.relu(self.conv3(cat1))
cat2 = torch.cat((x2, x3), 1)
x3 = self.ms(x3, mean1, std1)
x4 = F.relu(self.conv4(cat2))
x4 = self.ms(x4, mean2, std2)
cat3 = torch.cat((x1, x2, x3, x4), 1)
k = F.relu(self.conv5(cat3))
if k.size() != x.size():
raise Exception("k, haze image are different size!")
output = k * x - k + self.b
return F.relu(output)
class PONO(nn.Module):
def __init__(self, input_size=None, return_stats=False, affine=True, eps=1e-5):
super(PONO, self).__init__()
self.return_stats = return_stats
self.input_size = input_size
self.eps = eps
self.affine = affine
if affine:
self.beta = nn.Parameter(torch.zeros(1, 1, *input_size))
self.gamma = nn.Parameter(torch.ones(1, 1, *input_size))
else:
self.beta, self.gamma = None, None
def forward(self, x):
mean = x.mean(dim=1, keepdim=True)
std = (x.var(dim=1, keepdim=True) + self.eps).sqrt()
x = (x - mean) / std
if self.affine:
x = x * self.gamma + self.beta
return x, mean, std
class MS(nn.Module):
def __init__(self, beta=None, gamma=None):
super(MS, self).__init__()
self.gamma, self.beta = gamma, beta
def forward(self, x, beta=None, gamma=None):
beta = self.beta if beta is None else beta
gamma = self.gamma if gamma is None else gamma
if gamma is not None:
x.mul_(gamma)
if beta is not None:
x.add_(beta)
return x

View File

@@ -0,0 +1,138 @@
import os
import torch
import torch.backends.cudnn
import torch.nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torchvision
from torchvision import transforms
from torchvision.utils import make_grid
from tensorboardX import SummaryWriter
from utils import logger, weight_init
from config import get_config
from model import AODnet, AOD_pono_net
from data import HazeDataset
@logger
def load_data(cfg):
data_transform = transforms.Compose([
transforms.Resize([480, 640]),
transforms.ToTensor()
])
train_haze_dataset = HazeDataset(cfg.ori_data_path, cfg.haze_data_path, data_transform)
train_loader = torch.utils.data.DataLoader(train_haze_dataset, batch_size=cfg.batch_size, shuffle=True,
num_workers=cfg.num_workers, drop_last=True, pin_memory=True)
val_haze_dataset = HazeDataset(cfg.val_ori_data_path, cfg.val_haze_data_path, data_transform)
val_loader = torch.utils.data.DataLoader(val_haze_dataset, batch_size=cfg.val_batch_size, shuffle=False,
num_workers=cfg.num_workers, drop_last=True, pin_memory=True)
return train_loader, len(train_loader), val_loader, len(val_loader)
@logger
def save_model(epoch, path, net, optimizer, net_name):
if not os.path.exists(os.path.join(path, net_name)):
os.mkdir(os.path.join(path, net_name))
torch.save({'epoch': epoch, 'state_dict': net.state_dict(), 'optimizer': optimizer.state_dict()},
f=os.path.join(path, net_name, '{}_{}.pkl'.format('AOD', epoch)))
@logger
def load_network(device):
net = AOD_pono_net().to(device)
net.apply(weight_init)
return net
@logger
def load_optimizer(net, cfg):
optimizer = torch.optim.Adam(net.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
return optimizer
@logger
def loss_func(device):
criterion = torch.nn.MSELoss().to(device)
return criterion
@logger
def load_summaries(cfg):
summary = SummaryWriter(log_dir=os.path.join(cfg.log_dir, cfg.net_name), comment='')
return summary
def main(cfg):
# -------------------------------------------------------------------
# basic config
print(cfg)
if cfg.gpu > -1:
os.environ['CUDA_VISIBLE_DEVICES'] = str(cfg.gpu)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# -------------------------------------------------------------------
# load summaries
summary = load_summaries(cfg)
# -------------------------------------------------------------------
# load data
train_loader, train_number, val_loader, val_number = load_data(cfg)
# -------------------------------------------------------------------
# load loss
criterion = loss_func(device)
# -------------------------------------------------------------------
# load network
network = load_network(device)
# -------------------------------------------------------------------
# load optimizer
optimizer = load_optimizer(network, cfg)
# -------------------------------------------------------------------
# start train
print('Start train')
network.train()
for epoch in range(cfg.epochs):
for step, (ori_image, haze_image) in enumerate(train_loader):
count = epoch * train_number + (step + 1)
ori_image, haze_image = ori_image.to(device), haze_image.to(device)
dehaze_image = network(haze_image)
loss = criterion(dehaze_image, ori_image)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(network.parameters(), cfg.grad_clip_norm)
optimizer.step()
summary.add_scalar('loss', loss.item(), count)
if step % cfg.print_gap == 0:
summary.add_image('DeHaze_Images', make_grid(dehaze_image[:4].data, normalize=True, scale_each=True),
count)
summary.add_image('Haze_Images', make_grid(haze_image[:4].data, normalize=True, scale_each=True), count)
summary.add_image('Origin_Images', make_grid(ori_image[:4].data, normalize=True, scale_each=True),
count)
print('Epoch: {}/{} | Step: {}/{} | lr: {:.6f} | Loss: {:.6f}'
.format(epoch + 1, cfg.epochs, step + 1, train_number,
optimizer.param_groups[0]['lr'], loss.item()))
# -------------------------------------------------------------------
# start validation
print('Epoch: {}/{} | Validation Model Saving Images'.format(epoch + 1, cfg.epochs))
network.eval()
for step, (ori_image, haze_image) in enumerate(val_loader):
if step > 10: # only save image 10 times
break
ori_image, haze_image = ori_image.to(device), haze_image.to(device)
dehaze_image = network(haze_image)
torchvision.utils.save_image(
torchvision.utils.make_grid(torch.cat((haze_image, dehaze_image, ori_image), 0),
nrow=ori_image.shape[0]),
os.path.join(cfg.sample_output_folder, '{}_{}.jpg'.format(epoch + 1, step)))
network.train()
# -------------------------------------------------------------------
# save per epochs model
save_model(epoch, cfg.model_dir, network, optimizer, cfg.net_name)
# -------------------------------------------------------------------
# train finish
summary.close()
if __name__ == '__main__':
config_args, unparsed_args = get_config()
main(config_args)

View File

@@ -0,0 +1,18 @@
#!/usr/bin/env bash
python pono_train.py --epochs 10 \
--net_name aod-xavier \
--lr 1e-4 \
--use_gpu true \
--gpu 3 \
--ori_data_path data/images/ \
--haze_data_path data/data/ \
--val_ori_data_path data/images/ \
--val_haze_data_path data/val/ \
--num_workers 2 \
--batch_size 8 \
--val_batch_size 16 \
--print_gap 500 \
--model_dir ponomodels \
--log_dir ponologs \
--sample_output_folder ponosamples

Binary file not shown.

View File

@@ -0,0 +1,61 @@
#!/bin/bash
# 原图像位置
Dir_src_pics="./data/img"
Dir_result="./data/result"
Dir_ori_src_pics="/root/Dehaze/SRC_files/src"
mkdir -p $Dir_src_pics $Dir_result
# 进行绝对路径转化
Dir_src_pics=$(readlink -f "$Dir_src_pics")
Dir_result=$(readlink -f "$Dir_result")
if [ ! -d "$Dir_src_pics" ] && [ ! -d "$Dir_result" ]; then
echo "image、label都不存在程序退出"
echo -e "\033[31mori_image_directory\033[0m: $Dir_src_pics"
echo "$Dir_src_pics"
echo -e "\033[31mori_label_directory\033[0m: $Dir_result"
echo "$Dir_result"
exit 1
fi
PS3='All in one choice : '
applications=("Delete_src_pics" "Delete_generate_pics" "Copy_src_pics_to_test" "Run_test_program" "quit")
select fav in "${applications[@]}"; do
case $fav in
# 删除原始文件选项
"Delete_src_pics")
# 删除src文件
echo "Delete all src files in $Dir_src_pics"
rm $Dir_src_pics/*
;;
# 删除生成文件选项
"Delete_generate_pics")
# 删除result文件
echo "Delete all src files in $Dir_result"
rm $Dir_result/*
;;
# 复制待处理文件选项
"Copy_src_pics_to_test")
# 删除src文件
echo "Copy all src files in $Dir_ori_src_pics"
ln -s $Dir_ori_src_pics/* $Dir_src_pics
;;
# 运行程序
"Run_test_program")
source ~/miniconda/bin/activate Dehaze_DCP
python ./test/test.py $Dir_src_pics $Dir_result
;;
# 退出选项
"quit")
echo "User requested exit"
exit
;;
# 其他选项
*) echo "invalid option $REPLY";;
esac
done

Binary file not shown.

After

Width:  |  Height:  |  Size: 94 KiB

View File

@@ -0,0 +1,42 @@
# AOD-Net
ICCV 2017
AOD-Net is a light-weight but effective end-to-end dehazing neural network.
It is very easy and fast for you to train or test.
For test:
We have offered test.py, relevant prototxt and data.
You can just use 'python test.py' with GPU/CPU, you can obtain your result in data/result.
Good luck for your research.
# Improved AOD-Net with Positional Normalization (PONO)
We appreciate the [PyTorch implementation](https://github.com/TheFairBear/PyTorch-Image-Dehazing) of AOD-Net. Based on this code, we add simple [PONO](https://github.com/Boyiliee/PONO) into AOD-Net, which improves the performance efficiently.
Previous AOD-Net Results:
![](./AOD-Net_result.png)
For TestSet A, the PSNR increases from 19.69 to 20.38 dB, the SSIM increases from 0.8478 to 0.8587. For TestSetB, the PSNR increases from 21.54 to 21.67 dB, the SSIM increases from 0.9272 to 0.9285.
Please find in [AOD-Net with PONO](https://github.com/Boyiliee/AOD-Net/tree/master/AOD-Net%20with%20PONO) for details.
Bibtex:
```
@inproceedings{ICCV17a,
title={AOD-Net: All-in-One Dehazing Network},
author={Li, Boyi and Peng, Xiulian and Wang, Zhangyang and Xu, Ji-Zheng and Feng, Dan},
booktitle={Proceedings of the IEEE International Conference on Computer Vision},
year={2017}
}
@article{pono,
title={Positional Normalization},
author={Li, Boyi and Wu, Felix and Weinberger, Kilian Q. and Belongie, Serge},
journal={Advances in Neural Information Processing Systems},
year={2019}
}
```

View 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()

View 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"
}}

View File

@@ -0,0 +1,160 @@
import argparse
import os
import base64
import requests
import json
from PIL import Image
from io import BytesIO
def load_local_env():
"""
Load optional local credentials from config/baidu_api.env.
This file is ignored by git; see config/baidu_api.env.example.
"""
root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
env_path = os.path.join(root, 'config', 'baidu_api.env')
if not os.path.exists(env_path):
return
with open(env_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if not line or line.startswith('#') or '=' not in line:
continue
key, value = line.split('=', 1)
os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))
load_local_env()
API_KEY = os.environ.get("BAIDU_API_KEY", "")
SECRET_KEY = os.environ.get("BAIDU_SECRET_KEY", "")
# 输入和输出目录
INPUT_DIR = './去雾图像-北航合作'
OUTPUT_DIR = './Result_Baidu'
def get_access_token():
"""
使用 AKSK 生成鉴权签名Access Token
"""
if not API_KEY or not SECRET_KEY:
print("缺少 BAIDU_API_KEY 或 BAIDU_SECRET_KEY请设置环境变量或 config/baidu_api.env")
return None
url = "https://aip.baidubce.com/oauth/2.0/token"
params = {
"grant_type": "client_credentials",
"client_id": API_KEY,
"client_secret": SECRET_KEY
}
try:
response = requests.post(url, params=params)
return response.json().get("access_token")
except Exception as e:
print(f"获取 Access Token 失败: {e}")
return None
def process_image(file_path, access_token):
"""
读取图片调用百度API返回处理后的图片数据
"""
request_url = "https://aip.baidubce.com/rest/2.0/image-process/v1/dehaze"
request_url = request_url + "?access_token=" + access_token
try:
# 使用 PIL 读取图片
with Image.open(file_path) as img:
# 将 PIL Image 对象转换为字节流
img_buffer = BytesIO()
# 保持原格式保存到内存中
save_format = img.format if img.format else 'JPEG'
img.save(img_buffer, format=save_format)
img_bytes = img_buffer.getvalue()
# base64 编码
img_base64 = base64.b64encode(img_bytes)
params = {"image": img_base64}
headers = {'content-type': 'application/x-www-form-urlencoded'}
# 调用 API
response = requests.post(request_url, data=params, headers=headers)
if response.status_code == 200:
result = response.json()
# 检查是否有 image 字段
if "image" in result:
return base64.b64decode(result["image"])
else:
print(f"API 返回错误: {result}")
return None
else:
print(f"请求失败,状态码: {response.status_code}")
return None
except Exception as e:
print(f"处理图片 {file_path} 时出错: {e}")
return None
def main():
global INPUT_DIR, OUTPUT_DIR
parser = argparse.ArgumentParser(description="Run Baidu dehaze API on a folder of images.")
parser.add_argument("--input-dir", default=INPUT_DIR)
parser.add_argument("--output-dir", default=OUTPUT_DIR)
args = parser.parse_args()
INPUT_DIR = args.input_dir
OUTPUT_DIR = args.output_dir
# 1. 获取 Token
access_token = get_access_token()
if not access_token:
return
# 2. 确保输出目录存在
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
print(f"已创建输出目录: {OUTPUT_DIR}")
# 3. 遍历 ./Data 目录
if not os.path.exists(INPUT_DIR):
print(f"输入目录 {INPUT_DIR} 不存在")
return
supported_exts = ('.png', '.bmp', '.jpg', '.jpeg')
print("开始批量处理...")
files = os.listdir(INPUT_DIR)
for filename in files:
# 检查文件扩展名
if filename.lower().endswith(supported_exts):
input_path = os.path.join(INPUT_DIR, filename)
# V1
# output_path = os.path.join(OUTPUT_DIR, filename)
# V2
# 1. 分离文件名和扩展名
name, ext = os.path.splitext(filename)
# 2. 拼接新文件名
new_filename = f"{name}_result{ext}"
# 3. 生成最终输出路径
output_path = os.path.join(OUTPUT_DIR, new_filename)
print(f"正在处理: {filename} ...")
# 处理图片
processed_data = process_image(input_path, access_token)
# 保存结果
if processed_data:
with open(output_path, 'wb') as f:
f.write(processed_data)
print(f" 已保存至: {output_path}")
else:
print(f" 处理失败: {filename}")
print("批量处理完成。")
if __name__ == '__main__':
main()

View File

@@ -0,0 +1,4 @@
修改 1_Baidu_Dehaze.py 中 INPUT_DIR、Result_Baidu
python 1_Baidu_Dehaze.py
# 百度智能云控制台https://console.bce.baidu.com/ai-engine/imageprocess/overview/index

View File

@@ -0,0 +1,57 @@
#!/bin/bash
# 原图像位置
Dir_src_pics="./image/src"
Dir_dark="./image/dark"
Dir_result="./image/result"
Dir_trans="./image/trans"
Dir_ori_src_pics="/home/audience/Desktop/Dehaze/Dehaze/2025_11_23_SRC" # 修改这里
mkdir -p $Dir_src_pics $Dir_dark $Dir_result $Dir_trans
PS3='All in one choice : '
applications=("Delete_src_pics" "Delete_generate_pics" "Copy_src_pics" "Run_program" "quit")
select fav in "${applications[@]}"; do
case $fav in
# 删除原始文件选项
"Delete_src_pics")
# 删除src文件
echo "Delete all src files in $Dir_src_pics"
rm $Dir_src_pics/*
;;
# 删除生成文件选项
"Delete_generate_pics")
# 删除dark文件
echo "Delete all src files in $Dir_dark"
rm $Dir_dark/*
# 删除result文件
echo "Delete all src files in $Dir_result"
rm $Dir_result/*
# 删除trans文件
echo "Delete all src files in $Dir_trans"
rm $Dir_trans/*
;;
# 复制待处理文件选项
"Copy_src_pics")
# 删除src文件
echo "Copy all src files in $Dir_ori_src_pics"
ln -s $Dir_ori_src_pics/* $Dir_src_pics
;;
# 运行程序
"Run_program")
source ~/miniconda/bin/activate Dehaze_DCP
python dehaze.py
;;
# 退出选项
"quit")
echo "User requested exit"
exit
;;
# 其他选项
*) echo "invalid option $REPLY";;
esac
done

View File

@@ -0,0 +1,21 @@
The MIT License (MIT)
Copyright (c) 2016 WinCoder
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

View File

@@ -0,0 +1,18 @@
# single image dehaze
## Introduction
This program implement single image dehazing using dark channel prior.
## Compile Dependencies
- OpenCV
- Numpy
## Examples
<center>
<img src="./image/15.png" height = "400" alt="图片名称" />
<img src="./image/J.png" height = "400" alt="图片名称" />
</center>
## Algorithms
- Single Image Haze Removal Using Dark Channel Prior, Kaiming He, Jian Sun, and Xiaoou Tang", in CVPR 2009
- Guided Image Filtering, Kaiming He, Jian Sun, and Xiaoou Tang", in ECCV 2010.

View File

@@ -0,0 +1,150 @@
import cv2;
import math;
import os
import numpy as np;
def DarkChannel(im,sz):
b,g,r = cv2.split(im)
dc = cv2.min(cv2.min(r,g),b);
kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(sz,sz))
dark = cv2.erode(dc,kernel)
return dark
def AtmLight(im,dark):
[h,w] = im.shape[:2]
imsz = h*w
numpx = int(max(math.floor(imsz/1000),1))
darkvec = dark.reshape(imsz);
imvec = im.reshape(imsz,3);
indices = darkvec.argsort();
indices = indices[imsz-numpx::]
atmsum = np.zeros([1,3])
for ind in range(1,numpx):
atmsum = atmsum + imvec[indices[ind]]
A = atmsum / numpx;
return A
def TransmissionEstimate(im,A,sz):
omega = 0.95;
im3 = np.empty(im.shape,im.dtype);
for ind in range(0,3):
im3[:,:,ind] = im[:,:,ind]/A[0,ind]
transmission = 1 - omega*DarkChannel(im3,sz);
return transmission
def Guidedfilter(im,p,r,eps):
mean_I = cv2.boxFilter(im,cv2.CV_64F,(r,r));
mean_p = cv2.boxFilter(p, cv2.CV_64F,(r,r));
mean_Ip = cv2.boxFilter(im*p,cv2.CV_64F,(r,r));
cov_Ip = mean_Ip - mean_I*mean_p;
mean_II = cv2.boxFilter(im*im,cv2.CV_64F,(r,r));
var_I = mean_II - mean_I*mean_I;
a = cov_Ip/(var_I + eps);
b = mean_p - a*mean_I;
mean_a = cv2.boxFilter(a,cv2.CV_64F,(r,r));
mean_b = cv2.boxFilter(b,cv2.CV_64F,(r,r));
q = mean_a*im + mean_b;
return q;
def TransmissionRefine(im,et):
gray = cv2.cvtColor(im,cv2.COLOR_BGR2GRAY);
gray = np.float64(gray)/255;
r = 60;
eps = 0.0001;
t = Guidedfilter(gray,et,r,eps);
return t;
def Recover(im,t,A,tx = 0.1):
res = np.empty(im.shape,im.dtype);
t = cv2.max(t,tx);
for ind in range(0,3):
res[:,:,ind] = (im[:,:,ind]-A[0,ind])/t + A[0,ind]
return res
def getFileList(dir,Filelist, ext=None):
"""
获取文件夹及其子文件夹中文件列表
输入 dir文件夹根目录
输入 ext: 扩展名
返回: 文件路径列表
"""
newDir = dir
if os.path.isfile(dir):
if ext is None:
Filelist.append(dir)
else:
if ext in dir[-3:]:
Filelist.append(dir)
elif os.path.isdir(dir):
for s in os.listdir(dir):
newDir=os.path.join(dir,s)
getFileList(newDir, Filelist, ext)
return Filelist
if __name__ == '__main__':
import sys
sz = 10 # 窗口函数
tx = 0.2 # 传输图最小值
try:
local_dir = sys.argv[1]
except:
local_dir = './image/'
try:
sz = int(sys.argv[2])
except:
sz = sz
try:
tx = float(sys.argv[3])
except:
tx = tx
def nothing(*argv):
pass
# 检索文件
src_img_folder = os.path.join(local_dir, 'src')
imglist = getFileList(src_img_folder, [], '')
print('本次执行检索到 '+str(len(imglist))+' 张图像\n')
for imgpath in imglist:
imgname= os.path.splitext(os.path.basename(imgpath))[0]
# cv2.IMREAD_GRAYSCALE / cv2.IMREAD_COLOR 加载灰色 / 彩色图像
src = cv2.imread(imgpath, cv2.IMREAD_COLOR)
# 通道归一滑
I = src.astype('float64')/255;
# 暗通道图像
dark = DarkChannel(I,sz=sz);
#
A = AtmLight(I,dark);
te = TransmissionEstimate(I,A,sz = sz);
t = TransmissionRefine(src,te);
print(np.shape(src))
print(np.shape(dark))
print(np.shape(t))
J = Recover(I,t,A,tx=tx); # tx传输图的最小值用于避免过度曝光
dark_imgdir = os.path.join(local_dir, 'dark/')
trans_imgdir = os.path.join(local_dir, 'trans/')
result_imgdir = os.path.join(local_dir, 'result/')
cv2.imwrite(dark_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_dark.png",dark*255);
cv2.imwrite(trans_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_t.png", t*255);
cv2.imwrite(result_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_result.png",J*255);

57
DehazeNet/All_in_One.sh Normal file
View File

@@ -0,0 +1,57 @@
#!/bin/bash
# 原图像位置
Dir_src_pics="./img/src"
Dir_Trans_Refine="./img/Trans_Refine"
Dir_Trans_Esti="./img/Trans_Esti"
Dir_result="./img/result"
Dir_ori_src_pics="/root/Dehaze/SRC_files/src_1280_720"
mkdir -p $Dir_src_pics $Dir_Trans_Refine $Dir_Trans_Esti $Dir_result
PS3='All in one choice : '
applications=("Delete_src_pics" "Delete_generate_pics" "Copy_src_pics" "Run_program" "quit")
select fav in "${applications[@]}"; do
case $fav in
# 删除原始文件选项
"Delete_src_pics")
# 删除src文件
echo "Delete all src files in $Dir_src_pics"
rm $Dir_src_pics/*
;;
# 删除生成文件选项
"Delete_generate_pics")
# 删除dark文件
echo "Delete all src files in $Dir_Trans_Refine"
rm $Dir_Trans_Refine/*
# 删除result文件
echo "Delete all src files in $Dir_result"
rm $Dir_result/*
# 删除trans文件
echo "Delete all src files in $Dir_Trans_Esti"
rm $Dir_Trans_Esti/*
;;
# 复制待处理文件选项
"Copy_src_pics")
# 删除src文件
echo "Copy all src files in $Dir_ori_src_pics"
ln -s $Dir_ori_src_pics/* $Dir_src_pics
;;
# 运行程序
"Run_program")
source ~/miniconda/bin/activate Dehaze_DCP
python DehazeNet.py ./img
;;
# 退出选项
"quit")
echo "User requested exit"
exit
;;
# 其他选项
*) echo "invalid option $REPLY";;
esac
done

View File

@@ -0,0 +1,173 @@
name: "Dehaze_fullconv"
input: "data"
input_dim: 1
input_dim: 3
input_dim: {height_15}
input_dim: {width_15}
layer {{
name: "conv1"
type: "Convolution"
bottom: "data"
top: "conv1"
convolution_param {{
num_output: 20
kernel_size: 5
stride: 1
pad: 0
}}
}}
layer {{
name: "relu1"
type: "ReLU"
bottom: "conv1"
top: "conv1"
}}
layer {{
name: "reshape1"
type: "Reshape"
bottom: "conv1"
top: "reshape1"
reshape_param {{
shape {{
dim: 0
dim: 1
dim: 20
dim: -1
}}
}}
}}
layer {{
name: "pool1"
type: "Pooling"
bottom: "reshape1"
top: "pool1"
pooling_param {{
pool: MAX
kernel_w: 1
kernel_h: 5
stride_w: 1
stride_h: 5
}}
}}
layer {{
name: "reshape2"
type: "Reshape"
bottom: "pool1"
top: "reshape2"
reshape_param {{
shape {{
dim: 0
dim: 4
dim: {height_11}
dim: {width_11}
}}
}}
}}
layer {{
name: "conv2/1x1"
type: "Convolution"
bottom: "reshape2"
top: "conv2/1x1"
convolution_param {{
num_output: 16
kernel_size: 1
stride: 1
pad: 0
}}
}}
layer {{
name: "conv2/3x3"
type: "Convolution"
bottom: "reshape2"
top: "conv2/3x3"
convolution_param {{
num_output: 16
kernel_size: 3
stride: 1
pad: 1
}}
}}
layer {{
name: "conv2/5x5"
type: "Convolution"
bottom: "reshape2"
top: "conv2/5x5"
convolution_param {{
num_output: 16
kernel_size: 5
stride: 1
pad: 2
}}
}}
layer {{
name: "conv2/7x7"
type: "Convolution"
bottom: "reshape2"
top: "conv2/7x7"
convolution_param {{
num_output: 16
kernel_size: 7
stride: 1
pad: 3
}}
}}
layer {{
name: "conv2/output"
type: "Concat"
bottom: "conv2/1x1"
bottom: "conv2/3x3"
bottom: "conv2/5x5"
bottom: "conv2/7x7"
top: "conv2/output"
concat_param
{{
axis: 1
}}
}}
layer {{
name: "relu2"
type: "ReLU"
bottom: "conv2/output"
top: "conv2/output"
}}
layer {{
name: "pool2"
type: "Pooling"
bottom: "conv2/output"
top: "pool2"
pooling_param {{
pool: MAX
kernel_size: 8
stride: 1
}}
}}
layer {{
name: "ip1-conv"
type: "Convolution"
bottom: "pool2"
top: "ip1-conv"
convolution_param {{
num_output: 1
kernel_size: 5
}}
}}
layer {{
name: "drelu1"
type: "ReLU"
bottom: "ip1-conv"
top: "ip1-conv"
}}

Binary file not shown.

View File

@@ -0,0 +1,256 @@
name: "Dehaze"
input: "data"
input_dim: 1
input_dim: 3
input_dim: 16
input_dim: 16
layer {
name: "conv1"
type: "Convolution"
bottom: "data"
top: "conv1"
param {
lr_mult: 1
}
param {
lr_mult: 0.1
}
convolution_param {
num_output: 20
kernel_size: 5
stride: 1
pad: 0
weight_filler {
type: "gaussian"
std: 0.001
}
bias_filler {
type: "constant"
value: 0
}
}
}
layer {
name: "relu1"
type: "ReLU"
bottom: "conv1"
top: "conv1"
}
layer {
name: "reshape1"
type: "Reshape"
bottom: "conv1"
top: "reshape1"
reshape_param {
shape {
dim: 0
dim: 1
dim: 20
dim: -1
}
}
}
layer {
name: "pool1"
type: "Pooling"
bottom: "reshape1"
top: "pool1"
pooling_param {
pool: MAX
kernel_w: 1
kernel_h: 5
stride_w: 1
stride_h: 5
}
}
layer {
name: "reshape2"
type: "Reshape"
bottom: "pool1"
top: "reshape2"
reshape_param {
shape {
dim: 0
dim: 4
dim: 12
dim: 12
}
}
}
layer {
name: "conv2/1x1"
type: "Convolution"
bottom: "reshape2"
top: "conv2/1x1"
param {
lr_mult: 0.1
}
param {
lr_mult: 0.1
}
convolution_param {
num_output: 16
kernel_size: 1
stride: 1
pad: 0
weight_filler {
type: "gaussian"
std: 0.001
}
bias_filler {
type: "constant"
value: 0
}
}
}
layer {
name: "conv2/3x3"
type: "Convolution"
bottom: "reshape2"
top: "conv2/3x3"
param {
lr_mult: 0.1
}
param {
lr_mult: 0.1
}
convolution_param {
num_output: 16
kernel_size: 3
stride: 1
pad: 1
weight_filler {
type: "gaussian"
std: 0.001
}
bias_filler {
type: "constant"
value: 0
}
}
}
layer {
name: "conv2/5x5"
type: "Convolution"
bottom: "reshape2"
top: "conv2/5x5"
param {
lr_mult: 0.1
}
param {
lr_mult: 0.1
}
convolution_param {
num_output: 16
kernel_size: 5
stride: 1
pad: 2
weight_filler {
type: "gaussian"
std: 0.001
}
bias_filler {
type: "constant"
value: 0
}
}
}
layer {
name: "conv2/7x7"
type: "Convolution"
bottom: "reshape2"
top: "conv2/7x7"
param {
lr_mult: 0.1
}
param {
lr_mult: 0.1
}
convolution_param {
num_output: 16
kernel_size: 7
stride: 1
pad: 3
weight_filler {
type: "gaussian"
std: 0.001
}
bias_filler {
type: "constant"
value: 0
}
}
}
layer {
name: "conv2/output"
type: "Concat"
bottom: "conv2/1x1"
bottom: "conv2/3x3"
bottom: "conv2/5x5"
bottom: "conv2/7x7"
top: "conv2/output"
concat_param
{
axis: 1
}
}
layer {
name: "relu2"
type: "ReLU"
bottom: "conv2/output"
top: "conv2/output"
}
layer {
name: "pool2"
type: "Pooling"
bottom: "conv2/output"
top: "pool2"
pooling_param {
pool: MAX
kernel_size: 8
stride: 1
}
}
layer {
name: "ip1"
type: "InnerProduct"
bottom: "pool2"
top: "ip1"
param {
lr_mult: 1
}
param {
lr_mult: 2
}
inner_product_param {
num_output: 1
weight_filler {
type: "xavier"
}
bias_filler {
type: "constant"
}
}
}
layer {
name: "drelu1"
type: "ReLU"
bottom: "ip1"
top: "ip1"
}

208
DehazeNet/DehazeNet.py Normal file
View File

@@ -0,0 +1,208 @@
import sys,os
import caffe
import numpy as np
import cv2
import math
def EditFcnProto(templateFile, height, width):
with open(templateFile, 'r') as ft:
template = ft.read()
outFile = 'DehazeNetFcn.prototxt'
with open(outFile, 'w') as fd:
fd.write(template.format(height_15=height+15, width_15=width+15,
height_11=height+11, width_11=width+11))
def TransmissionEstimate(im_path, height, width):
caffe.set_mode_cpu()
# Define a safe tile size to prevent INT_MAX overflow (approx 512x512 is safe)
SAFE_TILE_SIZE = 512
# Use tiling if the image is larger than the safe size
if height > SAFE_TILE_SIZE or width > SAFE_TILE_SIZE:
print(f"Image size ({width}x{height}) is large. Using tiling to avoid memory overflow...")
# Determine effective tile size (cannot be larger than image)
tile_h = min(height, SAFE_TILE_SIZE)
tile_w = min(width, SAFE_TILE_SIZE)
# Generate prototxt for the TILE size, not the full image size
EditFcnProto('DehazeFcnTemplate.prototxt', tile_h, tile_w)
# Load networks
net = caffe.Net('DehazeNet.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
net_full_conv = caffe.Net('DehazeNetFcn.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
net_full_conv.params['ip1-conv'][0].data.flat = net.params['ip1'][0].data.flat
net_full_conv.params['ip1-conv'][1].data[...] = net.params['ip1'][1].data
# Load and pad image
im = caffe.io.load_image(im_path)
npad = ((7,8), (7,8), (0,0))
im_padded = np.pad(im, npad, 'symmetric')
transmission = np.zeros((height, width))
# Setup transformer for the tile size
transformers = caffe.io.Transformer({'data': net_full_conv.blobs['data'].data.shape})
transformers.set_transpose('data', (2,0,1))
transformers.set_channel_swap('data', (2,1,0))
# Process in tiles
for h in range(0, height, tile_h):
for w in range(0, width, tile_w):
# Calculate start/end to handle edges/overlap
# If we are at the end, shift back to ensure we feed a full tile
h_start = min(h, height - tile_h)
w_start = min(w, width - tile_w)
# Extract patch from PADDED image
# Network expects input size of (Tile + 15), so we slice accordingly
patch = im_padded[h_start : h_start + tile_h + 15, w_start : w_start + tile_w + 15, :]
# Forward pass
out = net_full_conv.forward_all(data=np.array([transformers.preprocess('data', patch-0.2)]))
# Reshape output
block_trans = np.reshape(out['ip1-conv'], (tile_h, tile_w))
# Assign to result buffer
transmission[h_start : h_start + tile_h, w_start : w_start + tile_w] = block_trans
return transmission
else:
# Original logic for small images
EditFcnProto('DehazeFcnTemplate.prototxt', height, width)
net = caffe.Net('DehazeNet.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
net_full_conv = caffe.Net('DehazeNetFcn.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
net_full_conv.params['ip1-conv'][0].data.flat = net.params['ip1'][0].data.flat
net_full_conv.params['ip1-conv'][1].data[...] = net.params['ip1'][1].data
im = caffe.io.load_image(im_path)
npad = ((7,8), (7,8), (0,0))
im = np.pad(im, npad, 'symmetric')
transformers = caffe.io.Transformer({'data': net_full_conv.blobs['data'].data.shape})
transformers.set_transpose('data', (2,0,1))
transformers.set_channel_swap('data', (2,1,0))
out = net_full_conv.forward_all(data=np.array([transformers.preprocess('data', im-0.2)]))
transmission = np.reshape(out['ip1-conv'], (height,width))
return transmission
def DarkChannel(im,sz):
b,g,r = cv2.split(im)
dc = cv2.min(cv2.min(r,g),b)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(sz,sz))
dark = cv2.erode(dc,kernel)
return dark
def AtmLight(im,dark):
[h,w] = im.shape[:2]
imsz = h*w
numpx = int(max(math.floor(imsz/1000),1))
darkvec = dark.reshape(imsz,1)
imvec = im.reshape(imsz,3)
indices = darkvec.argsort()
indices = indices[imsz-numpx::]
atmsum = np.zeros([1,3])
for ind in range(1,numpx):
atmsum = atmsum + imvec[indices[ind]]
A = atmsum / numpx
return A
def Guidedfilter(im,p,r,eps):
mean_I = cv2.boxFilter(im,cv2.CV_64F,(r,r))
mean_p = cv2.boxFilter(p, cv2.CV_64F,(r,r))
mean_Ip = cv2.boxFilter(im*p,cv2.CV_64F,(r,r))
cov_Ip = mean_Ip - mean_I*mean_p
mean_II = cv2.boxFilter(im*im,cv2.CV_64F,(r,r))
var_I = mean_II - mean_I*mean_I
a = cov_Ip/(var_I + eps)
b = mean_p - a*mean_I
mean_a = cv2.boxFilter(a,cv2.CV_64F,(r,r))
mean_b = cv2.boxFilter(b,cv2.CV_64F,(r,r))
q = mean_a*im + mean_b
return q
def TransmissionRefine(im,et):
gray = cv2.cvtColor(im,cv2.COLOR_BGR2GRAY)
gray = np.float64(gray)/255
r = 60
eps = 0.0001
t = Guidedfilter(gray,et,r,eps)
return t
def Recover(im,t,A,tx = 0.1):
res = np.empty(im.shape,im.dtype)
t = cv2.max(t,tx)
for ind in range(0,3):
res[:,:,ind] = (im[:,:,ind]-A[0,ind])/t + A[0,ind]
return res
def getFileList(dir,Filelist, ext=None):
"""
获取文件夹及其子文件夹中文件列表
输入 dir文件夹根目录
输入 ext: 扩展名
返回: 文件路径列表
"""
newDir = dir
if os.path.isfile(dir):
if ext is None:
Filelist.append(dir)
else:
if ext in dir[-3:]:
Filelist.append(dir)
elif os.path.isdir(dir):
for s in os.listdir(dir):
newDir=os.path.join(dir,s)
getFileList(newDir, Filelist, ext)
return Filelist
if __name__ == '__main__':
if not len(sys.argv) == 2:
print ('Usage: python DeHazeNet.py haze_img_path')
exit()
else:
im_path = sys.argv[1]
# 检索文件
src_img_folder = os.path.join(im_path, 'src')
imglist = getFileList(src_img_folder, [], '')
print('本次执行检索到 '+str(len(imglist))+' 张图像\n')
for img_path in imglist:
imgname= os.path.splitext(os.path.basename(img_path))[0]
src = cv2.imread(img_path)
height = src.shape[0]
width = src.shape[1]
# Note: EditFcnProto is also called inside TransmissionEstimate if tiling is used
# We call it here for initialization but it may be overwritten.
templateFile = 'DehazeFcnTemplate.prototxt'
EditFcnProto(templateFile, height, width)
print("-"*5, ' 完成EditFcnProto ',"-"*5)
I = src/255.0
dark = DarkChannel(I,15)
A = AtmLight(I,dark)
te = TransmissionEstimate(img_path, height, width)
t = TransmissionRefine(src,te)
J = Recover(I,t,A,0.1)
print("Finsh All the operation")
Trans_Esti_imgdir = os.path.join(im_path, 'Trans_Esti/')
if not os.path.exists(Trans_Esti_imgdir): os.makedirs(Trans_Esti_imgdir)
print(Trans_Esti_imgdir + imgname + "_Trans_Esti.png")
cv2.imwrite(Trans_Esti_imgdir + imgname + "_Trans_Esti.png",te*255);
Trans_Refine_imgdir = os.path.join(im_path, 'Trans_Refine/')
if not os.path.exists(Trans_Refine_imgdir): os.makedirs(Trans_Refine_imgdir)
print(Trans_Refine_imgdir + imgname + "_Trans_Refine.png")
cv2.imwrite(Trans_Refine_imgdir + imgname + "_Trans_Refine.png",t*255);
result_imgdir = os.path.join(im_path, 'result/')
if not os.path.exists(result_imgdir): os.makedirs(result_imgdir)
print(result_imgdir + imgname + "_result.png")
cv2.imwrite(result_imgdir + imgname + "_result.png",J*255);

31
DehazeNet/readme.md Normal file
View File

@@ -0,0 +1,31 @@
## Reimplement
## *DehazeNet: An End-to-End System for Single Image Haze Removal*
Bolun Cai, Xiangmin Xu, Kui Jia, Chunmei Qing, Dacheng Tao
## Requirement
> * caffe
> * opencv2
## Usage:
simply type
```shell
python DehazeNet.py image_path
```
## Demo:
![canon](img/canon.jpg)
![canon_Dehaze](img/canon_Dehaze.jpg)
![cones](img/cones.jpg)
![cones_Dehaze](img/cones_Dehaze.jpg)
## Site:
@article{cai2016dehazenet,
title={Dehazenet: An end-to-end system for single image haze removal},
author={Cai, Bolun and Xu, Xiangmin and Jia, Kui and Qing, Chunmei and Tao, Dacheng},
journal={IEEE Transactions on Image Processing},
volume={25},
number={11},
pages={5187--5198},
year={2016},
publisher={IEEE}
}

49
GCANet/All_in_One.sh Normal file
View File

@@ -0,0 +1,49 @@
#!/bin/bash
# 原图像位置
Dir_src_pics="./imgs/src"
Dir_result="./imgs/result"
Dir_ori_src_pics="/root/Dehaze/SRC_files/src"
mkdir -p $Dir_src_pics $Dir_result
PS3='All in one choice : '
applications=("Delete_src_pics" "Delete_generate_pics" "Copy_src_pics" "Run_program" "quit")
select fav in "${applications[@]}"; do
case $fav in
# 删除原始文件选项
"Delete_src_pics")
# 删除src文件
echo "Delete all src files in $Dir_src_pics"
rm $Dir_src_pics/*
;;
# 删除生成文件选项
"Delete_generate_pics")
# 删除result文件
echo "Delete all src files in $Dir_result"
rm $Dir_result/*
;;
# 复制待处理文件选项
"Copy_src_pics")
# 删除src文件
echo "Copy all src files in $Dir_ori_src_pics"
ln -s $Dir_ori_src_pics/* $Dir_src_pics
;;
# 运行程序
"Run_program")
source ~/miniconda/bin/activate Dehaze_GCANet
python test.py --task dehaze --gpu_id 0 --indir ./imgs/src --outdir ./imgs/result
;;
# 退出选项
"quit")
echo "User requested exit"
exit
;;
# 其他选项
*) echo "invalid option $REPLY";;
esac
done

102
GCANet/GCANet.py Normal file
View File

@@ -0,0 +1,102 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class ShareSepConv(nn.Module):
def __init__(self, kernel_size):
super(ShareSepConv, self).__init__()
assert kernel_size % 2 == 1, 'kernel size should be odd'
self.padding = (kernel_size - 1)//2
weight_tensor = torch.zeros(1, 1, kernel_size, kernel_size)
weight_tensor[0, 0, (kernel_size-1)//2, (kernel_size-1)//2] = 1
self.weight = nn.Parameter(weight_tensor)
self.kernel_size = kernel_size
def forward(self, x):
inc = x.size(1)
expand_weight = self.weight.expand(inc, 1, self.kernel_size, self.kernel_size).contiguous()
return F.conv2d(x, expand_weight,
None, 1, self.padding, 1, inc)
class SmoothDilatedResidualBlock(nn.Module):
def __init__(self, channel_num, dilation=1, group=1):
super(SmoothDilatedResidualBlock, self).__init__()
self.pre_conv1 = ShareSepConv(dilation*2-1)
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
self.pre_conv2 = ShareSepConv(dilation*2-1)
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
def forward(self, x):
y = F.relu(self.norm1(self.conv1(self.pre_conv1(x))))
y = self.norm2(self.conv2(self.pre_conv2(y)))
return F.relu(x+y)
class ResidualBlock(nn.Module):
def __init__(self, channel_num, dilation=1, group=1):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
def forward(self, x):
y = F.relu(self.norm1(self.conv1(x)))
y = self.norm2(self.conv2(y))
return F.relu(x+y)
class GCANet(nn.Module):
def __init__(self, in_c=4, out_c=3, only_residual=True):
super(GCANet, self).__init__()
self.conv1 = nn.Conv2d(in_c, 64, 3, 1, 1, bias=False)
self.norm1 = nn.InstanceNorm2d(64, affine=True)
self.conv2 = nn.Conv2d(64, 64, 3, 1, 1, bias=False)
self.norm2 = nn.InstanceNorm2d(64, affine=True)
self.conv3 = nn.Conv2d(64, 64, 3, 2, 1, bias=False)
self.norm3 = nn.InstanceNorm2d(64, affine=True)
self.res1 = SmoothDilatedResidualBlock(64, dilation=2)
self.res2 = SmoothDilatedResidualBlock(64, dilation=2)
self.res3 = SmoothDilatedResidualBlock(64, dilation=2)
self.res4 = SmoothDilatedResidualBlock(64, dilation=4)
self.res5 = SmoothDilatedResidualBlock(64, dilation=4)
self.res6 = SmoothDilatedResidualBlock(64, dilation=4)
self.res7 = ResidualBlock(64, dilation=1)
self.gate = nn.Conv2d(64 * 3, 3, 3, 1, 1, bias=True)
self.deconv3 = nn.ConvTranspose2d(64, 64, 4, 2, 1)
self.norm4 = nn.InstanceNorm2d(64, affine=True)
self.deconv2 = nn.Conv2d(64, 64, 3, 1, 1)
self.norm5 = nn.InstanceNorm2d(64, affine=True)
self.deconv1 = nn.Conv2d(64, out_c, 1)
self.only_residual = only_residual
def forward(self, x):
y = F.relu(self.norm1(self.conv1(x)))
y = F.relu(self.norm2(self.conv2(y)))
y1 = F.relu(self.norm3(self.conv3(y)))
y = self.res1(y1)
y = self.res2(y)
y = self.res3(y)
y2 = self.res4(y)
y = self.res5(y2)
y = self.res6(y)
y3 = self.res7(y)
gates = self.gate(torch.cat((y1, y2, y3), dim=1))
gated_y = y1 * gates[:, [0], :, :] + y2 * gates[:, [1], :, :] + y3 * gates[:, [2], :, :]
y = F.relu(self.norm4(self.deconv3(gated_y)))
y = F.relu(self.norm5(self.deconv2(y)))
if self.only_residual:
y = self.deconv1(y)
else:
y = F.relu(self.deconv1(y))
return y

View File

@@ -0,0 +1,102 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class ShareSepConv(nn.Module):
def __init__(self, kernel_size):
super(ShareSepConv, self).__init__()
assert kernel_size % 2 == 1, 'kernel size should be odd'
self.padding = (kernel_size - 1)//2
weight_tensor = torch.zeros(1, 1, kernel_size, kernel_size)
weight_tensor[0, 0, (kernel_size-1)//2, (kernel_size-1)//2] = 1
self.weight = nn.Parameter(weight_tensor)
self.kernel_size = kernel_size
def forward(self, x):
inc = x.size(1)
expand_weight = self.weight.expand(inc, 1, self.kernel_size, self.kernel_size).contiguous()
return F.conv2d(x, expand_weight,
None, 1, self.padding, 1, inc)
class SmoothDilatedResidualBlock(nn.Module):
def __init__(self, channel_num, dilation=1, group=1):
super(SmoothDilatedResidualBlock, self).__init__()
self.pre_conv1 = ShareSepConv(dilation*2-1)
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
self.pre_conv2 = ShareSepConv(dilation*2-1)
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
def forward(self, x):
y = F.relu(self.norm1(self.conv1(self.pre_conv1(x))))
y = self.norm2(self.conv2(self.pre_conv2(y)))
return F.relu(x+y)
class ResidualBlock(nn.Module):
def __init__(self, channel_num, dilation=1, group=1):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
def forward(self, x):
y = F.relu(self.norm1(self.conv1(x)))
y = self.norm2(self.conv2(y))
return F.relu(x+y)
class GCANet(nn.Module):
def __init__(self, in_c=4, out_c=3, only_residual=True):
super(GCANet, self).__init__()
self.conv1 = nn.Conv2d(in_c, 64, 3, 1, 1, bias=False)
self.norm1 = nn.InstanceNorm2d(64, affine=True)
self.conv2 = nn.Conv2d(64, 64, 3, 1, 1, bias=False)
self.norm2 = nn.InstanceNorm2d(64, affine=True)
self.conv3 = nn.Conv2d(64, 64, 3, 2, 1, bias=False)
self.norm3 = nn.InstanceNorm2d(64, affine=True)
self.res1 = SmoothDilatedResidualBlock(64, dilation=2)
self.res2 = SmoothDilatedResidualBlock(64, dilation=2)
self.res3 = SmoothDilatedResidualBlock(64, dilation=2)
self.res4 = SmoothDilatedResidualBlock(64, dilation=4)
self.res5 = SmoothDilatedResidualBlock(64, dilation=4)
self.res6 = SmoothDilatedResidualBlock(64, dilation=4)
self.res7 = ResidualBlock(64, dilation=1)
self.gate = nn.Conv2d(64 * 3, 3, 3, 1, 1, bias=True)
self.deconv3 = nn.ConvTranspose2d(64, 64, 4, 2, 1)
self.norm4 = nn.InstanceNorm2d(64, affine=True)
self.deconv2 = nn.Conv2d(64, 64, 3, 1, 1)
self.norm5 = nn.InstanceNorm2d(64, affine=True)
self.deconv1 = nn.Conv2d(64, out_c, 1)
self.only_residual = only_residual
def forward(self, x):
y = F.relu(self.norm1(self.conv1(x)))
y = F.relu(self.norm2(self.conv2(y)))
y1 = F.relu(self.norm3(self.conv3(y)))
y = self.res1(y1)
y = self.res2(y)
y = self.res3(y)
y2 = self.res4(y)
y = self.res5(y2)
y = self.res6(y)
y3 = self.res7(y)
gates = self.gate(torch.cat((y1, y2, y3), dim=1))
gated_y = y1 * gates[:, [0], :, :] + y2 * gates[:, [1], :, :] + y3 * gates[:, [2], :, :]
y = F.relu(self.norm4(self.deconv3(gated_y)))
y = F.relu(self.norm5(self.deconv2(y)))
if self.only_residual:
y = self.deconv1(y)
else:
y = F.relu(self.deconv1(y))
return y

View File

@@ -0,0 +1,104 @@
import os
import bisect
import threading
import torch
import numpy as np
import numpy.random as random
from PIL import Image
from torch.utils.data import Dataset
from folder_loader import FolderLoader
import torchvision.transforms as transforms
from utils import batch_edge_compute
def pil_loader(img_path):
return Image.open(img_path).convert("RGB")
class ImagePairPrefixFolder(Dataset):
def __init__(self, input_folder, gt_folder, max_img_size=0, size_unit=1, force_rgb=False):
super(ImagePairPrefixFolder, self).__init__()
self.gt_loader = FolderLoader(gt_folder)
# build the map from image name to index
self.gt_map = dict()
for idx, img_name in enumerate(self.gt_loader.img_names):
self.gt_map[os.path.splitext(img_name)[0].split('_')[0]] = idx
self.input_loader = FolderLoader(input_folder)
assert all([os.path.splitext(x)[0].split('_')[0] in self.gt_map for x in self.input_loader.img_names]), \
'cannot find corresponding gt names'
self.input_folder = input_folder
self.gt_folder = gt_folder
self.max_img_size = max_img_size
self.size_unit = size_unit
self.force_rgb = force_rgb
def __getitem__(self, index):
input_name, input_img = self.input_loader[index]
input_basename = os.path.splitext(input_name)[0].split('_')[0]
gt_idx = self.gt_map[input_basename]
gt_name, gt_img = self.gt_loader[gt_idx]
if self.force_rgb:
input_img = input_img.convert('RGB')
gt_img = gt_img.convert('RGB')
im_w, im_h = input_img.size
gt_w, gt_h = gt_img.size
assert im_w==gt_w and im_h==gt_h, 'input image and gt image size not match'
im_w, im_h = input_img.size
if 0 < self.max_img_size < max(im_w, im_h):
if im_w < im_h:
out_h = int(self.max_img_size) // self.size_unit * self.size_unit
out_w = int(im_w / im_h * out_h) // self.size_unit * self.size_unit
else:
out_w = int(self.max_img_size) // self.size_unit * self.size_unit
out_h = int(im_h / im_w * out_w) // self.size_unit * self.size_unit
else:
out_w = im_w // self.size_unit * self.size_unit
out_h = im_h // self.size_unit * self.size_unit
if im_w != out_w or im_h != out_h:
input_img = input_img.resize((out_w, out_h), Image.BILINEAR)
gt_img = gt_img.resize((out_w, out_h), Image.BILINEAR)
im_w, im_h = input_img.size
input_img = np.array(input_img).astype('float')
gt_img = np.array(gt_img).astype('float')
if len(input_img.shape) == 2:
input_img = input_img[:, :, np.newaxis]
if len(gt_img.shape) == 2:
gt_img = gt_img[:, :, np.newaxis]
return {'input_img': input_img, 'gt_img': gt_img, 'input_h': im_h, "input_w": im_w}
def get_input_info(self, index):
image_name = os.path.splitext(self.input_loader.img_names[index])[0]
return self.input_loader, image_name
def __len__(self):
return len(self.input_loader)
def var_custom_collate(batch):
min_h, min_w = 10000, 10000
for item in batch:
min_h = min(min_h, item['input_h'])
min_w = min(min_w, item['input_w'])
inc = 1 if len(batch[0]['input_img'].shape)==2 else batch[0]['input_img'].shape[2]
batch_input_images = torch.Tensor(len(batch), inc, min_h, min_w)
batch_gt_images = torch.Tensor(len(batch), inc, min_h, min_w)
for idx, item in enumerate(batch):
off_y = 0 if item['input_h']==min_h else random.randint(0, item['input_h'] - min_h)
off_x = 0 if item['input_w']==min_w else random.randint(0, item['input_w'] - min_w)
crop_input_img = item['input_img'][off_y:off_y + min_h, off_x:off_x + min_w, :]
crop_gt_img = item['gt_img'][off_y:off_y + min_h, off_x:off_x + min_w, :]
batch_input_images[idx] = torch.from_numpy(crop_input_img.transpose((2, 0, 1))) - 128
batch_gt_images[idx] = torch.from_numpy(crop_gt_img.transpose((2, 0, 1)))
batch_input_edges = batch_edge_compute(batch_input_images) - 128
return batch_input_images, batch_input_edges, batch_gt_images

View File

@@ -0,0 +1,19 @@
import io
import os
import utils
import struct
from PIL import Image
class FolderLoader(object):
def __init__(self, fold_path):
super(FolderLoader, self).__init__()
self.fold_path = fold_path
self.img_paths = utils.make_dataset(self.fold_path)
self.img_names = [os.path.basename(x) for x in self.img_paths]
def __getitem__(self, index):
img = Image.open(self.img_paths[index])#.convert('RGB')
return self.img_names[index], img
def __len__(self):
return len(self.img_names)

View File

@@ -0,0 +1,66 @@
import os
import argparse
import numpy as np
from PIL import Image
import torch
from torch.autograd import Variable
from utils import make_dataset, edge_compute
parser = argparse.ArgumentParser()
parser.add_argument('--network', default='GCANet')
parser.add_argument('--task', default='dehaze', help='dehaze | derain')
parser.add_argument('--gpu_id', type=int, default=0)
parser.add_argument('--indir', default='examples/')
parser.add_argument('--outdir', default='output')
opt = parser.parse_args()
assert opt.task in ['dehaze', 'derain']
## forget to regress the residue for deraining by mistake,
## which should be able to produce better results
opt.only_residual = opt.task == 'dehaze'
opt.model = 'models/wacv_gcanet_%s.pth' % opt.task
opt.use_cuda = opt.gpu_id >= 0
if not os.path.exists(opt.outdir):
os.makedirs(opt.outdir)
test_img_paths = make_dataset(opt.indir)
if opt.network == 'GCANet':
from GCANet import GCANet
net = GCANet(in_c=4, out_c=3, only_residual=opt.only_residual)
else:
print('network structure %s not supported' % opt.network)
raise ValueError
if opt.use_cuda:
torch.cuda.set_device(opt.gpu_id)
net.cuda()
else:
net.float()
net.load_state_dict(torch.load(opt.model, map_location='cpu'))
net.eval()
for img_path in test_img_paths:
img = Image.open(img_path).convert('RGB')
im_w, im_h = img.size
if im_w % 4 != 0 or im_h % 4 != 0:
img = img.resize((int(im_w // 4 * 4), int(im_h // 4 * 4)))
img = np.array(img).astype('float')
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
edge_data = edge_compute(img_data)
in_data = torch.cat((img_data, edge_data), dim=0).unsqueeze(0) - 128
in_data = in_data.cuda() if opt.use_cuda else in_data.float()
with torch.no_grad():
pred = net(Variable(in_data))
if opt.only_residual:
out_img_data = (pred.data[0].cpu().float() + img_data).round().clamp(0, 255)
else:
out_img_data = pred.data[0].cpu().float().round().clamp(0, 255)
out_img = Image.fromarray(out_img_data.numpy().astype(np.uint8).transpose(1, 2, 0))
out_img.save(os.path.join(opt.outdir, os.path.splitext(os.path.basename(img_path))[0] + '_%s.png' % opt.task))

View File

@@ -0,0 +1,41 @@
import numpy as np
import os
import ntpath
import time
import utils
from scipy.misc import imresize
from tensorboardX import SummaryWriter
class TFVisualizer():
def __init__(self, opt):
self.tf_visualizer = SummaryWriter(os.path.join(opt.logDir, opt.name))
self.opt = opt
self.saved = False
self.ncols = 4
self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
with open(self.log_name, "a") as log_file:
now = time.strftime("%c")
log_file.write('================ Training Loss (%s) ================\n' % now)
def reset(self):
self.saved = False
# |visuals|: dictionary of images to display or save
def display_current_results(self, visuals, iter_mark, epoch, save_result):
for label, image in visuals.items():
img_gid = utils.tensor2imgrid(image)
self.tf_visualizer.add_image(label, img_gid, iter_mark)
# losses: dictionary of error labels and values
def plot_current_losses(self, iter_mark, losses):
# for label, loss in losses.items():
# self.tf_visualizer.add_scalar(label, loss, iter_mark)
self.tf_visualizer.add_scalars('training loss', losses, iter_mark)
def print_logs(self, message):
print(message)
with open(self.log_name, "a") as log_file:
log_file.write('%s\n' % message)

View File

@@ -0,0 +1,228 @@
import os
import datetime
import argparse
import numpy as np
import torch
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from ImagePairPrefixFolder import ImagePairPrefixFolder, var_custom_collate
from utils import MovingAvg
from tf_visualizer import TFVisualizer
parser = argparse.ArgumentParser()
parser.add_argument('--network', default='GCANet')
parser.add_argument('--name', default='default_exp')
parser.add_argument('--gpu_ids', default='0')
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--lr_step', type=int, default=40)
parser.add_argument('--lr_gamma', type=float, default=0.1)
parser.add_argument('--weight_decay', type=float, default=0.0005)
parser.add_argument('--checkpoints_dir', default='checkpoint')
parser.add_argument('--logDir', default='tblogdir')
parser.add_argument('--resume_dir', default='')
parser.add_argument('--resume_epoch', type=int, default=0)
parser.add_argument('--save_epoch', type=int, default=5)
parser.add_argument('--save_latest_freq', type=int, default=5000)
parser.add_argument('--test_epoch', type=int, default=5)
parser.add_argument('--test_max_size', type=int, default=1080)
parser.add_argument('--size_unit', type=int, default=8)
parser.add_argument('--print_iter', type=int, default=100)
parser.add_argument('--input_folder', default='')
parser.add_argument('--gt_folder', default='')
parser.add_argument('--test_input_folder', default='')
parser.add_argument('--test_gt_folder', default='')
parser.add_argument('--num_workers', type=int, default=16)
parser.add_argument('--batch_size', type=int, default=4)
parser.add_argument('--only_residual', action='store_true', help='regress residual rather than image')
parser.add_argument('--loss_func', default='l2', help='l2|l1')
parser.add_argument('--inc', type=int, default=3)
parser.add_argument('--outc', type=int, default=3)
parser.add_argument('--force_rgb', action='store_true')
parser.add_argument('--no_edge', action='store_true')
opt = parser.parse_args()
opt.input_folder = os.path.expanduser(opt.input_folder)
opt.gt_folder = os.path.expanduser(opt.gt_folder)
opt.test_input_folder = os.path.expanduser(opt.test_input_folder)
opt.test_gt_folder = os.path.expanduser(opt.test_gt_folder)
if not os.path.exists(os.path.join(opt.checkpoints_dir, opt.name)):
os.makedirs(os.path.join(opt.checkpoints_dir, opt.name))
opt.resume_dir = opt.resume_dir if opt.resume_dir != '' else os.path.join(opt.checkpoints_dir, opt.name)
visualizer = TFVisualizer(opt)
### Log out
with open(os.path.realpath(__file__), 'r') as fid:
visualizer.print_logs(fid.read())
## print argument
for key, val in vars(opt).items():
visualizer.print_logs('%s: %s' % (key, val))
opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(',')]
assert all(0 <= x <= torch.cuda.device_count() for x in opt.gpu_ids), 'gpu id should ' \
'be 0~{0}'.format(torch.cuda.device_count())
torch.cuda.set_device(opt.gpu_ids[0])
train_dataset = ImagePairPrefixFolder(opt.input_folder, opt.gt_folder, size_unit=opt.size_unit, force_rgb=opt.force_rgb)
train_dataloader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True,
collate_fn=var_custom_collate, pin_memory=True,
num_workers=opt.num_workers)
opt.do_test = opt.test_gt_folder != ''
if opt.do_test:
test_dataset = ImagePairPrefixFolder(opt.test_input_folder, opt.test_gt_folder,
max_img_size=opt.test_max_size, size_unit=opt.size_unit, force_rgb=opt.force_rgb)
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False,
collate_fn=var_custom_collate, pin_memory=True,
num_workers=1)
total_inc = opt.inc if opt.no_edge else opt.inc + 1
if opt.network == 'GCANet':
from GCANet import GCANet
net = GCANet(in_c=total_inc, out_c=3, only_residual=opt.only_residual)
else:
print('network structure %s not supported' % opt.network)
raise ValueError
if opt.loss_func == 'l2':
loss_crit = torch.nn.MSELoss()
elif opt.loss_func == 'l1':
loss_crit = torch.nn.SmoothL1Loss()
else:
print('loss_func %s not supported' % opt.loss_func)
raise ValueError
pnsr_crit = torch.nn.MSELoss()
if len(opt.gpu_ids) > 0:
net.cuda()
if len(opt.gpu_ids) > 1:
net = torch.nn.DataParallel(net)
loss_crit = loss_crit.cuda()
pnsr_crit = pnsr_crit.cuda()
optimizer = optim.Adam(net.parameters(), lr=opt.lr)
step_optim_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt.lr_step, gamma=opt.lr_gamma)
loss_avg = MovingAvg(pool_size=50)
start_epoch = 0
total_iter = 0
if os.path.exists(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth')):
print('resuming from latest.pth')
latest_info = torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth'))
start_epoch = latest_info['epoch']
total_iter = latest_info['total_iter']
if isinstance(net, torch.nn.DataParallel):
net.module.load_state_dict(latest_info['net_state'])
else:
net.load_state_dict(latest_info['net_state'])
optimizer.load_state_dict(latest_info['optim_state'])
if opt.resume_epoch > 0:
start_epoch = opt.resume_epoch
total_iter = opt.resume_epoch * len(train_dataloader)
resume_path = os.path.join(opt.resume_epoch, 'net_epoch_%d.pth') % opt.resume_epoch
print('resume from : %s' % resume_path)
assert os.path.exists(resume_path), 'cannot find the resume model: %s ' % resume_path
if isinstance(net, torch.nn.DataParallel):
net.module.load_state_dict(torch.load(resume_path))
else:
net.load_state_dict(torch.load(resume_path))
for epoch in range(start_epoch, opt.epochs):
visualizer.print_logs("Start to train epoch %d" % epoch)
net.train()
for iter, data in enumerate(train_dataloader):
total_iter += 1
optimizer.zero_grad()
step_optim_scheduler.step(epoch)
batch_input_img, batch_input_edge, batch_gt = data
if len(opt.gpu_ids) > 0:
batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda()
if opt.no_edge:
batch_input = batch_input_img
else:
batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1)
batch_input_v = Variable(batch_input)
if opt.only_residual:
batch_gt_v = Variable(batch_gt - (batch_input_img+128))
else:
batch_gt_v = Variable(batch_gt)
pred = net(batch_input_v)
loss = loss_crit(pred, batch_gt_v)
avg_loss = loss_avg.set_curr_val(loss.data)
loss.backward()
optimizer.step()
if iter % opt.print_iter == 0:
visualizer.plot_current_losses(total_iter, { 'loss': loss})
visualizer.print_logs('%s Step[%d/%d], lr: %f, mv_avg_loss: %f, loss: %f' %
(str(datetime.datetime.now()).split(' ')[1], iter, len(train_dataloader),
step_optim_scheduler.get_lr()[0], avg_loss, loss))
if total_iter % opt.save_latest_freq == 0:
latest_info = {'total_iter': total_iter,
'epoch': epoch,
'optim_state': optimizer.state_dict()}
if len(opt.gpu_ids) > 1:
latest_info['net_state'] = net.module.state_dict()
else:
latest_info['net_state'] = net.state_dict()
print('save lastest model.')
torch.save(latest_info, os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth'))
if (epoch+1) % opt.save_epoch == 0 :
visualizer.print_logs('saving model for epoch %d' % epoch)
if len(opt.gpu_ids) > 1:
torch.save(net.module.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch+1)))
else:
torch.save(net.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch + 1)))
if opt.do_test:
avg_psnr = 0
task_cnt = 0
net.eval()
with torch.no_grad():
for iter, data in enumerate(test_dataloader):
batch_input_img, batch_input_edge, batch_gt = data
if len(opt.gpu_ids) > 0:
batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda()
if opt.no_edge:
batch_input = batch_input_img
else:
batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1)
batch_input_v = Variable(batch_input)
batch_gt_v = Variable(batch_gt)
pred = net(batch_input_v)
if opt.only_residual:
loss = pnsr_crit(pred+Variable(batch_input_img+128), batch_gt_v)
else:
loss = pnsr_crit(pred, batch_gt_v)
avg_psnr += 10 * np.log10(255 * 255 / loss.item())
task_cnt += 1
visualizer.print_logs('Testing for epoch: %d' % epoch)
visualizer.print_logs('Average test PNSR is %f for %d images' % (avg_psnr/task_cnt, task_cnt))

View File

@@ -0,0 +1,198 @@
import os
import torch
import torch
import numpy as np
from PIL import Image
import os
from scipy import signal
from torchvision.utils import make_grid
IMG_EXTENSIONS = [
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def make_dataset(dir):
images = []
assert os.path.isdir(dir), '%s is not a valid directory' % dir
for root, _, fnames in sorted(os.walk(dir)):
for fname in fnames:
if is_image_file(fname):
path = os.path.join(root, fname)
images.append(path)
return images
def edge_compute(x):
x_diffx = torch.abs(x[:,:,1:] - x[:,:,:-1])
x_diffy = torch.abs(x[:,1:,:] - x[:,:-1,:])
y = x.new(x.size())
y.fill_(0)
y[:,:,1:] += x_diffx
y[:,:,:-1] += x_diffx
y[:,1:,:] += x_diffy
y[:,:-1,:] += x_diffy
y = torch.sum(y,0,keepdim=True)/3
y /= 4
return y
def batch_edge_compute(x):
x_diffx = torch.abs(x[:,:,:,1:] - x[:,:,:,:-1])
x_diffy = torch.abs(x[:,:,1:,:] - x[:,:,:-1,:])
y = x.new(x.size())
y.fill_(0)
y[:,:,:,1:] += x_diffx
y[:,:,:,:-1] += x_diffx
y[:,:,1:,:] += x_diffy
y[:,:,:-1,:] += x_diffy
y = torch.sum(y,1,keepdim=True)/3
y /= 4
return y
# Converts a Tensor into an image array (numpy)
# |imtype|: the desired type of the converted numpy array
def tensor2im(input_image, imtype=np.uint8):
if isinstance(input_image, torch.Tensor):
image_tensor = input_image.data
else:
return input_image
image_numpy = image_tensor[0].cpu().float().numpy()
if image_numpy.shape[0] == 1:
image_numpy = np.tile(image_numpy, (3, 1, 1))
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
image_numpy = image_numpy.clip(0, 255)
return image_numpy.astype(imtype)
def tensor2imgrid(input_image):
im_grid = make_grid(input_image[:4, ...], nrow=2, normalize=True, range=(-128, 128))
return im_grid
# ndarr = im_grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
# im = Image.fromarray(ndarr)
# return im
def diagnose_network(net, name='network'):
mean = 0.0
count = 0
for param in net.parameters():
if param.grad is not None:
mean += torch.mean(torch.abs(param.grad.data))
count += 1
if count > 0:
mean = mean / count
print(name)
print(mean)
def save_image(image_numpy, image_path):
image_pil = Image.fromarray(image_numpy)
image_pil.save(image_path)
def print_numpy(x, val=True, shp=False):
x = x.astype(np.float64)
if shp:
print('shape,', x.shape)
if val:
x = x.flatten()
print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
def mkdirs(paths):
if isinstance(paths, list) and not isinstance(paths, str):
for path in paths:
mkdir(path)
else:
mkdir(paths)
def mkdir(path):
if not os.path.exists(path):
os.makedirs(path)
def fspecial_gauss(size, sigma):
"""Function to mimic the 'fspecial' gaussian MATLAB function
"""
x, y = np.mgrid[-size // 2 + 1:size // 2 + 1, -size // 2 + 1:size // 2 + 1]
g = np.exp(-((x ** 2 + y ** 2) / (2.0 * sigma ** 2)))
return g / g.sum()
def filter2(x, kernel, mode='same'):
return signal.convolve2d(x, np.rot90(kernel, 2), mode=mode)
def ssim(img1, img2, cs_map=False):
"""Return the Structural Similarity Map corresponding to input images img1
and img2 (images are assumed to be uint8)
This function attempts to mimic precisely the functionality of ssim.m a
MATLAB provided by the author's of SSIM
https://ece.uwaterloo.ca/~z70wang/research/ssim/ssim_index.m
"""
img1 = img1.astype(np.float64)
img2 = img2.astype(np.float64)
size = 11
sigma = 1.5
window = fspecial_gauss(size, sigma)
K1 = 0.01
K2 = 0.03
L = 255 # bitdepth of image
C1 = (K1 * L) ** 2
C2 = (K2 * L) ** 2
mu1 = filter2(img1, window, mode='valid')
mu2 = filter2(img2, window, mode='valid')
mu1_sq = mu1 * mu1
mu2_sq = mu2 * mu2
mu1_mu2 = mu1 * mu2
sigma1_sq = filter2(img1 * img1, window, mode='valid') - mu1_sq
sigma2_sq = filter2(img2 * img2, window, mode='valid') - mu2_sq
sigma12 = filter2(img1 * img2, window, mode='valid') - mu1_mu2
if cs_map:
return np.mean(np.mean((((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
(sigma1_sq + sigma2_sq + C2)),
(2.0 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2))))
else:
return np.mean(np.mean(((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
(sigma1_sq + sigma2_sq + C2))))
class MovingAvg(object):
def __init__(self, pool_size=100):
from queue import Queue
self.pool = Queue(maxsize=pool_size)
self.sum = 0
self.curr_pool_size = 0
self.pool_size = pool_size
def set_curr_val(self, val):
if not self.pool.full():
self.curr_pool_size += 1
self.pool.put_nowait(val)
else:
last_first_val = self.pool.get_nowait()
self.pool.put_nowait(val)
self.sum -= last_first_val
self.sum += val
return self.sum / self.curr_pool_size
def reset(self):
from queue import Queue
self.pool = Queue(maxsize=self.pool_size)
self.sum = 0
self.curr_pool_size = 0

46
GCANet/README.md Normal file
View File

@@ -0,0 +1,46 @@
Gated Context Aggregation Network for Image Dehazing and Deraining
=======
![image](imgs/net_arch.png)
This is the implementation of our WACV 2019 paper *"Gated Context Aggregation Network for Image Dehazing and Deraining"* by [Dongdong Chen](<http://www.dongdongchen.bid/>), [Mingming He](<https://github.com/hmmlillian>), [Qingnan Fan](<https://fqnchina.github.io/>), *et al.*
In this paper, we propose a new end-to-end gated context aggregation network GCANet for image dehazing, in which the smoothed dilated convolution is used to avoid the gridding artifacts and a gated subnetwork is applied to fuse the features of different levels. Experiments show that GCANet can obtain much better performance than all the previous state-of-the-art image dehazing methods both qualitatively and quantitatively
![image](imgs/dehaze_visual.png)
We further apply our proposed GCANet to the image deraining task, which also outperforms previous state-of-the-art image deraining methods and demonstrates its generality.
![image](imgs/derain_visual.png)
## Getting Started
This paper is implemented with Pytorch framework.
Demo
----
Directly put all your test images under one directory. Then run:
```bash
python test.py --task [dehaze | derain] --gpu_id [gpu_id] --indir [input directory] --outdir [output directory]
```
For training, please download the training code from <https://drive.google.com/file/d/1T7X1HYztbz6S75vTRNtREgGEOI269KDk/view?usp=sharing>
Cite
----
You can use our codes for research purpose only. And please cite our paper when you use our codes.
```
@article{chen2018gated,
title={Gated Context Aggregation Network for Image Dehazing and Deraining},
author={Chen, Dongdong and He, Mingming and Fan, Qingnan and Liao, Jing and Zhang, Liheng and Hou, Dongdong and Yuan, Lu and Hua, Gang},
journal={WACV 2019},
year={2018}
}
```
Contact
-------
If you find any bugs or have any ideas of optimizing these codes, please contact me via cddlyf [at] gmail [dot] com

Binary file not shown.

Binary file not shown.

67
GCANet/test.py Normal file
View File

@@ -0,0 +1,67 @@
import os
import argparse
import numpy as np
from PIL import Image
import torch
from torch.autograd import Variable
from utils import make_dataset, edge_compute
parser = argparse.ArgumentParser()
parser.add_argument('--network', default='GCANet')
parser.add_argument('--task', default='dehaze', help='dehaze | derain')
parser.add_argument('--gpu_id', type=int, default=0)
parser.add_argument('--indir', default='examples/')
parser.add_argument('--outdir', default='output')
opt = parser.parse_args()
assert opt.task in ['dehaze', 'derain']
## forget to regress the residue for deraining by mistake,
## which should be able to produce better results
opt.only_residual = opt.task == 'dehaze'
opt.model = 'models/wacv_gcanet_%s.pth' % opt.task
opt.use_cuda = opt.gpu_id >= 0
if not os.path.exists(opt.outdir):
os.makedirs(opt.outdir)
test_img_paths = make_dataset(opt.indir)
if opt.network == 'GCANet':
from GCANet import GCANet
net = GCANet(in_c=4, out_c=3, only_residual=opt.only_residual)
else:
print('network structure %s not supported' % opt.network)
raise ValueError
if opt.use_cuda:
torch.cuda.set_device(opt.gpu_id)
net.cuda()
else:
net.float()
net.load_state_dict(torch.load(opt.model, map_location='cpu'))
net.eval()
for img_path in test_img_paths:
img = Image.open(img_path).convert('RGB')
im_w, im_h = img.size
if im_w % 4 != 0 or im_h % 4 != 0:
img = img.resize((int(im_w // 4 * 4), int(im_h // 4 * 4)))
img = np.array(img).astype('float')
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
edge_data = edge_compute(img_data)
in_data = torch.cat((img_data, edge_data), dim=0).unsqueeze(0) - 128
in_data = in_data.cuda() if opt.use_cuda else in_data.float()
with torch.no_grad():
pred = net(Variable(in_data))
if opt.only_residual:
out_img_data = (pred.data[0].cpu().float() + img_data).round().clamp(0, 255)
else:
out_img_data = pred.data[0].cpu().float().round().clamp(0, 255)
out_img = Image.fromarray(out_img_data.numpy().astype(np.uint8).transpose(1, 2, 0))
print("-"*5,"图片存储在:",os.path.join(opt.outdir, os.path.splitext(os.path.basename(img_path))[0] + '_%s.png' % opt.task))
out_img.save(os.path.join(opt.outdir, os.path.splitext(os.path.basename(img_path))[0] + '_%s.png' % opt.task))

39
GCANet/utils.py Normal file
View File

@@ -0,0 +1,39 @@
import os
import torch
IMG_EXTENSIONS = [
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def make_dataset(dir):
images = []
assert os.path.isdir(dir), '%s is not a valid directory' % dir
for root, _, fnames in sorted(os.walk(dir)):
for fname in fnames:
if is_image_file(fname):
path = os.path.join(root, fname)
images.append(path)
return images
def edge_compute(x):
x_diffx = torch.abs(x[:,:,1:] - x[:,:,:-1])
x_diffy = torch.abs(x[:,1:,:] - x[:,:-1,:])
y = x.new(x.size())
y.fill_(0)
y[:,:,1:] += x_diffx
y[:,:,:-1] += x_diffx
y[:,1:,:] += x_diffy
y[:,:-1,:] += x_diffy
y = torch.sum(y,0,keepdim=True)/3
y /= 4
return y

37
README.md Normal file
View File

@@ -0,0 +1,37 @@
# Dehaze
多模型图像去雾对比与后处理网页工具。
## 快速启动
```bash
./run_dehaze_web.sh
```
访问:
```text
http://192.168.3.11:7860/
```
详细说明见 [使用手册.md](使用手册.md)。
## 已整合能力
- AOD
- Baidu_API
- DCP
- DehazeNet
- GCANet
- RefineDNet
- 多种 HSV/SV 后处理方式
## 验证
```bash
python scripts/verify_all.py --images 1.png
```
运行结果保存在 `web_results/`,该目录不提交到 git。
当前已在本机验证 `1.png``2.jpg` 的六个模型和四种后处理流程均可运行。

50
RefineDNet/All_in_One.sh Normal file
View File

@@ -0,0 +1,50 @@
#!/bin/bash
# 原图像位置
Dir_pic_root="./datasets/quick_test"
Dir_src_pics="./datasets/quick_test/src"
Dir_result="./datasets/quick_test/result"
Dir_ori_src_pics="/root/Dehaze/SRC_files/src"
mkdir -p $Dir_pic_root $Dir_src_pics $Dir_result
PS3='All in one choice : '
applications=("Delete_src_pics" "Delete_generate_pics" "Copy_src_pics" "Run_program" "quit")
select fav in "${applications[@]}"; do
case $fav in
# 删除原始文件选项
"Delete_src_pics")
# 删除src文件
echo "Delete all src files in $Dir_src_pics"
rm $Dir_src_pics/*
;;
# 删除生成文件选项
"Delete_generate_pics")
# 删除result文件
echo "Delete all src files in $Dir_result"
rm $Dir_result/*
;;
# 复制待处理文件选项
"Copy_src_pics")
# 删除src文件
echo "Copy all src files in $Dir_ori_src_pics"
ln -s $Dir_ori_src_pics/* $Dir_src_pics
;;
# 运行程序
"Run_program")
source ~/miniconda/bin/activate Dehaze_GCANet
python quick_test.py --dataroot $Dir_pic_root --dataset_mode single --name refined_DCP_outdoor --model refined_DCP --phase src --preprocess none --save_image --method_name refined_DCP_outdoor_ep_60 --epoch 60
;;
# 退出选项
"quit")
echo "User requested exit"
exit
;;
# 其他选项
*) echo "invalid option $REPLY";;
esac
done

103
RefineDNet/README.md Normal file
View File

@@ -0,0 +1,103 @@
# RefineDNet for dehazing
RefineDNet is a two-stage dehazing framework which can be weakly supervised using real-world unpaired images.
That is, the training set never requires paired hazy and haze-free images coming from the same scene.
In the first stage, it adopts DCP to restore visibility of the input hazy image.
In the second stage, it improves the realness of preliminary results from the first stage via CNNs.
RefineDNet is outlined in the following figure, and more details can be found in the [paper](https://doi.org/10.1109/TIP.2021.3060873) (or [this link](https://sse.tongji.edu.cn/linzhang/files/RefineDNet_TIP.pdf)) titled as _RefineDNet: A Weakly Supervised Refinement Framework for Single Image Dehazing._ (Early Access in Trans. Image Process.)
![framework](https://github.com/xiaofeng94/RefineDNet-for-dehazing/blob/master/datasets/figures/framework_github.jpg)
# Our Environment
- Ubuntu 16.06
- Python (>= 3.5)
- PyTorch (>= 1.1.0) with CUDA 9.0
- torchvision (>=0.3.0)
- numpy (>= 1.17.0)
# Testing
## Download the pretrained models.
1. Get the model on [Google drive](https://drive.google.com/file/d/1NIm-o01AOdjGn3kvsVA57TEn6jYNKGr4/view?usp=sharing) or [BaiduYun Disk](https://pan.baidu.com/s/1pqy-Ka9b9xVaeumdNSZAWQ) (Key: bswu). It's trained on RESIDE-unpaired.
2. Create a folder named `checkpoints`, and unzip `refined_DCP_outdoor.zip` in `./checkpoints`.
Now, your directory tree should look like
```
<RefineDNet_root>
├── checkpoints
│ ├── refined_DCP_outdoor
│ │ ├── 60_net_D.pth
│ │ ├── 60_net_Refiner_J.pth
│ │ ├── 60_net_Refiner_T.pth
│ │ └── test_opt.txt
│ ...
...
```
## Quick test on real-world images
1. Download the pretrained model on RESIDE-unpaired (see above).
2. Run the following command from <RefineDNet_root>.
```
python quick_test.py --dataroot ./datasets/quick_test --dataset_mode single --name refined_DCP_outdoor --model refined_DCP --phase test --preprocess none --save_image --method_name refined_DCP_outdoor_ep_60 --epoch 60
```
The results will be saved in the folder `<RefineDNet_root>/datatsets/quick_test/refined_DCP_outdoor_ep_60`.
## Test on BeDDE
1. Download the pretrained model on BeDDE.
2. Run the following command from `<RefineDNet_root>`.
```
python test_BeDDE.py --dataroot <BeDDE_root> --dataset_mode simple_bedde --bedde_list ./datasets/BeDDE/bedde_list.txt --name refined_DCP_outdoor --model refined_DCP --phase test --preprocess none --save_image --method_name refined_DCP_outdoor_ep_60 --epoch 60
```
The results will be saved in the folder `<BeDDE_root>/<city_name>/refined_DCP_outdoor_ep_60`.
# Training
## Train RefineDNet on RESIDE-unpaired
1. Download RESIDE-unpaired on [Google drive](https://drive.google.com/file/d/1SjQwESy8nwVO7pC3JRW7vXvJ6Qqk6Et4/view?usp=sharing) or [BaiduYun Disk](https://pan.baidu.com/s/1pqy-Ka9b9xVaeumdNSZAWQ) (Key: bswu). Unzip `RESIDE-unpaired.zip` in the folder <RefineDNet_root>/datasets.
Your directory tree should look like
```
<RefineDNet_root>
├── datasets
│ ├── BeDDE
│ ├── RESIDE-unpaired
│ │ ├── trainA
│ │ └── trainB
│ ...
...
```
2. Open visdom by `python -m visdom.server`
3. Run the following command from `<RefineDNet_root>`.
```
python train.py --dataroot ./datasets/RESIDE-unpaired --dataset_mode unpaired --model refined_DCP --name refined_DCP_outdoor --niter 30 --niter_decay 60 --lr_decay_iters 10 --preprocess scale_min_and_crop --load_size 300 --crop_size 256 --num_threads 8 --save_epoch_freq 3
```
## Train RefineDNet on ITS (from RESIDE-standard)
1. Download ITS [here](https://sites.google.com/view/reside-dehaze-datasets/reside-standard?authuser=0). Unzip hazy.zip and clear.zip into `<RefineDNet_root>/datasets/ITS`.
2. Rename the hazy image folder as `trainA` and the clear image folder as `trainB`.
Then, your directory tree should look like
```
<RefineDNet_root>
├── datasets
│ ├── BeDDE
│ ├── ITS
│ │ ├── trainA
│ │ └── trainB
│ ...
...
```
3. Open visdom by `python -m visdom.server`
4. Run the following command from `<RefineDNet_root>`.
```
python train.py --dataroot ./datasets/ITS --dataset_mode unpaired --model refined_DCP --name refined_DCP_indoor --niter 30 --niter_decay 60 --lr_decay_iters 5 --preprocess scale_width_and_crop --load_size 372 --crop_size 256 --num_threads 8 --save_epoch_freq 1
```
# Results
Some dehazing samples from BeDDE and the Internet produced by various methods.
![dehazing samples](https://github.com/xiaofeng94/RefineDNet-for-dehazing/blob/master/datasets/figures/outdoor_com_github.jpg)
# Useful links
1. [RESIDE dataset](https://sites.google.com/view/reside-dehaze-datasets/reside-standard?authuser=0)
2. [BeDDE dataset](https://github.com/xiaofeng94/BeDDE-for-defogging)
3. This code is based on [CycleGAN](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix)

View File

@@ -0,0 +1,44 @@
----------------- Options ---------------
aspect_ratio: 1.0
batch_size: 1
checkpoints_dir: ./checkpoints
crop_size: 256
dataroot: ./datasets/quick_test [default: ./datasets/ITS_v2]
dataset_mode: single [default: unaligned]
direction: AtoB
display_winsize: 256
epoch: 60 [default: latest]
eval: False
gpu_ids: 0
init_gain: 0.02
init_type: normal
input_nc: 3
isTrain: False [default: None]
load_iter: 0 [default: 0]
load_size: 256
max_dataset_size: inf
method_name: refined_DCP_outdoor_ep_60 [default: Mine]
model: refined_DCP [default: test]
n_layers_D: 3
name: refined_DCP_outdoor [default: experiment_name]
ndf: 64
netD: basic
netG: resnet_9blocks
netR_J: resnet_9blocks
netR_T: unet_trans_256
ngf: 64
no_dropout: True
no_flip: False
norm: instance
ntest: inf
num_test: 50
num_threads: 4
output_nc: 3
phase: src [default: test]
preprocess: none [default: resize_and_crop]
results_dir: ./results/
save_image: True [default: False]
serial_batches: False
suffix:
verbose: False
----------------- End -------------------

View File

@@ -0,0 +1,44 @@
----------------- Options ---------------
aspect_ratio: 1.0
batch_size: 1
checkpoints_dir: ./checkpoints
crop_size: 256
dataroot: /home/wkmgc/Desktop/Dehaze/web_results/2_08fda024/work/RefineDNet/dataset [default: ./datasets/ITS_v2]
dataset_mode: single [default: unaligned]
direction: AtoB
display_winsize: 256
epoch: 60 [default: latest]
eval: False
gpu_ids: 0
init_gain: 0.02
init_type: normal
input_nc: 3
isTrain: False [default: None]
load_iter: 0 [default: 0]
load_size: 256
max_dataset_size: inf
method_name: RefineDNet [default: Mine]
model: refined_DCP [default: test]
n_layers_D: 3
name: refined_DCP_outdoor [default: experiment_name]
ndf: 64
netD: basic
netG: resnet_9blocks
netR_J: resnet_9blocks
netR_T: unet_trans_256
ngf: 64
no_dropout: True
no_flip: False
norm: instance
ntest: inf
num_test: 50
num_threads: 4
output_nc: 3
phase: test
preprocess: none [default: resize_and_crop]
results_dir: ./results/
save_image: True [default: False]
serial_batches: False
suffix:
verbose: False
----------------- End -------------------

View File

@@ -0,0 +1,93 @@
"""This package includes all the modules related to data loading and preprocessing
To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
You need to implement four functions:
-- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
-- <__len__>: return the size of dataset.
-- <__getitem__>: get a data point from data loader.
-- <modify_commandline_options>: (optionally) add dataset-specific options and set default options.
Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
See our template dataset class 'template_dataset.py' for more details.
"""
import importlib
import torch.utils.data
from data.base_dataset import BaseDataset
def find_dataset_using_name(dataset_name):
"""Import the module "data/[dataset_name]_dataset.py".
In the file, the class called DatasetNameDataset() will
be instantiated. It has to be a subclass of BaseDataset,
and it is case-insensitive.
"""
dataset_filename = "data." + dataset_name + "_dataset"
datasetlib = importlib.import_module(dataset_filename)
dataset = None
target_dataset_name = dataset_name.replace('_', '') + 'dataset'
for name, cls in datasetlib.__dict__.items():
if name.lower() == target_dataset_name.lower() \
and issubclass(cls, BaseDataset):
dataset = cls
if dataset is None:
raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
return dataset
def get_option_setter(dataset_name):
"""Return the static method <modify_commandline_options> of the dataset class."""
dataset_class = find_dataset_using_name(dataset_name)
return dataset_class.modify_commandline_options
def create_dataset(opt):
"""Create a dataset given the option.
This function wraps the class CustomDatasetDataLoader.
This is the main interface between this package and 'train.py'/'test.py'
Example:
>>> from data import create_dataset
>>> dataset = create_dataset(opt)
"""
data_loader = CustomDatasetDataLoader(opt)
dataset = data_loader.load_data()
return dataset
class CustomDatasetDataLoader():
"""Wrapper class of Dataset class that performs multi-threaded data loading"""
def __init__(self, opt):
"""Initialize this class
Step 1: create a dataset instance given the name [dataset_mode]
Step 2: create a multi-threaded data loader.
"""
self.opt = opt
dataset_class = find_dataset_using_name(opt.dataset_mode)
self.dataset = dataset_class(opt)
print("dataset [%s] was created" % type(self.dataset).__name__)
self.dataloader = torch.utils.data.DataLoader(
self.dataset,
batch_size=opt.batch_size,
shuffle=not opt.serial_batches,
num_workers=int(opt.num_threads))
def load_data(self):
return self
def __len__(self):
"""Return the number of data in the dataset"""
return min(len(self.dataset), self.opt.max_dataset_size)
def __iter__(self):
"""Return a batch of data"""
for i, data in enumerate(self.dataloader):
if i * self.opt.batch_size >= self.opt.max_dataset_size:
break
yield data

View File

@@ -0,0 +1,60 @@
import os.path
from data.base_dataset import BaseDataset, get_params, get_transform
from data.image_folder import make_dataset
from PIL import Image
class AlignedDataset(BaseDataset):
"""A dataset class for paired image dataset.
It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.
During test time, you need to prepare a directory '/path/to/data/test'.
"""
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_AB = os.path.join(opt.dataroot, opt.phase) # get the image directory
self.AB_paths = sorted(make_dataset(self.dir_AB, opt.max_dataset_size)) # get image paths
assert(self.opt.load_size >= self.opt.crop_size) # crop_size should be smaller than the size of loaded image
self.input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
self.output_nc = self.opt.input_nc if self.opt.direction == 'BtoA' else self.opt.output_nc
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index - - a random integer for data indexing
Returns a dictionary that contains A, B, A_paths and B_paths
A (tensor) - - an image in the input domain
B (tensor) - - its corresponding image in the target domain
A_paths (str) - - image paths
B_paths (str) - - image paths (same as A_paths)
"""
# read a image given a random integer index
AB_path = self.AB_paths[index]
AB = Image.open(AB_path).convert('RGB')
# split AB image into A and B
w, h = AB.size
w2 = int(w / 2)
A = AB.crop((0, 0, w2, h))
B = AB.crop((w2, 0, w, h))
# apply the same transform to both A and B
transform_params = get_params(self.opt, A.size)
A_transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))
B_transform = get_transform(self.opt, transform_params, grayscale=(self.output_nc == 1))
A = A_transform(A)
B = B_transform(B)
return {'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path}
def __len__(self):
"""Return the total number of images in the dataset."""
return len(self.AB_paths)

View File

@@ -0,0 +1,177 @@
"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
"""
import random
import numpy as np
import torch.utils.data as data
from PIL import Image
import torchvision.transforms as transforms
from abc import ABC, abstractmethod
class BaseDataset(data.Dataset, ABC):
"""This class is an abstract base class (ABC) for datasets.
To create a subclass, you need to implement the following four functions:
-- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
-- <__len__>: return the size of dataset.
-- <__getitem__>: get a data point.
-- <modify_commandline_options>: (optionally) add dataset-specific options and set default options.
"""
def __init__(self, opt):
"""Initialize the class; save the options in the class
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
self.opt = opt
self.root = opt.dataroot
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
return parser
@abstractmethod
def __len__(self):
"""Return the total number of images in the dataset."""
return 0
@abstractmethod
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index - - a random integer for data indexing
Returns:
a dictionary of data with their names. It ususally contains the data itself and its metadata information.
"""
pass
def get_params(opt, size):
w, h = size
new_h = h
new_w = w
if opt.preprocess == 'resize_and_crop':
new_h = new_w = opt.load_size
elif opt.preprocess == 'scale_width_and_crop':
new_w = opt.load_size
new_h = opt.load_size * h // w
elif opt.preprocess == 'scale_min_and_crop':
if w <= h:
new_w = opt.load_size
new_h = opt.load_size * h // w
else:
new_w = opt.load_size * w // h
new_h = opt.load_size
x = random.randint(0, np.maximum(0, new_w - opt.crop_size))
y = random.randint(0, np.maximum(0, new_h - opt.crop_size))
flip = random.random() > 0.5
return {'crop_pos': (x, y), 'flip': flip}
def get_transform(opt, params=None, grayscale=False, method=Image.BICUBIC, convert=True):
transform_list = []
if grayscale:
transform_list.append(transforms.Grayscale(1))
if 'resize' in opt.preprocess:
osize = [opt.load_size, opt.load_size]
transform_list.append(transforms.Resize(osize, method))
elif 'scale_width' in opt.preprocess:
transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method)))
elif 'scale_min' in opt.preprocess:
transform_list.append(transforms.Lambda(lambda img: __scale_min(img, opt.load_size, method)))
if 'crop' in opt.preprocess:
if params is None:
transform_list.append(transforms.RandomCrop(opt.crop_size))
else:
transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size)))
if opt.preprocess == 'none':
transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))
if not opt.no_flip:
if params is None:
transform_list.append(transforms.RandomHorizontalFlip())
elif params['flip']:
transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
if convert:
transform_list += [transforms.ToTensor()]
if grayscale:
transform_list += [transforms.Normalize((0.5,), (0.5,))]
else:
transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
return transforms.Compose(transform_list)
def __make_power_2(img, base, method=Image.BICUBIC):
ow, oh = img.size
h = int(round(oh / base) * base)
w = int(round(ow / base) * base)
if (h == oh) and (w == ow):
return img
__print_size_warning(ow, oh, w, h)
return img.resize((w, h), method)
def __scale_width(img, target_width, method=Image.BICUBIC):
ow, oh = img.size
if (ow == target_width):
return img
w = target_width
h = int(target_width * oh / ow)
return img.resize((w, h), method)
def __scale_min(img, target_min, method=Image.BICUBIC):
ow, oh = img.size
if ow <= oh:
return __scale_width(img, target_min, method)
else:
if (oh == target_min):
return img
w = int(target_min * ow/oh)
h = target_min
return img.resize((w, h), method)
def __crop(img, pos, size):
ow, oh = img.size
x1, y1 = pos
tw = th = size
if (ow > tw or oh > th):
return img.crop((x1, y1, x1 + tw, y1 + th))
return img
def __flip(img, flip):
if flip:
return img.transpose(Image.FLIP_LEFT_RIGHT)
return img
def __print_size_warning(ow, oh, w, h):
"""Print warning information about image size(only print once)"""
if not hasattr(__print_size_warning, 'has_printed'):
print("The image size needs to be a multiple of 4. "
"The loaded image size was (%d, %d), so it was adjusted to "
"(%d, %d). This adjustment will be done to all images "
"whose sizes are not multiples of 4" % (ow, oh, w, h))
__print_size_warning.has_printed = True

View File

@@ -0,0 +1,68 @@
import os.path
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from skimage import color # require skimage
from PIL import Image
import numpy as np
import torchvision.transforms as transforms
class ColorizationDataset(BaseDataset):
"""This dataset class can load a set of natural images in RGB, and convert RGB format into (L, ab) pairs in Lab color space.
This dataset is required by pix2pix-based colorization model ('--model colorization')
"""
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
By default, the number of channels for input image is 1 (L) and
the nubmer of channels for output image is 2 (ab). The direction is from A to B
"""
parser.set_defaults(input_nc=1, output_nc=2, direction='AtoB')
return parser
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir = os.path.join(opt.dataroot, opt.phase)
self.AB_paths = sorted(make_dataset(self.dir, opt.max_dataset_size))
assert(opt.input_nc == 1 and opt.output_nc == 2 and opt.direction == 'AtoB')
self.transform = get_transform(self.opt, convert=False)
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index - - a random integer for data indexing
Returns a dictionary that contains A, B, A_paths and B_paths
A (tensor) - - the L channel of an image
B (tensor) - - the ab channels of the same image
A_paths (str) - - image paths
B_paths (str) - - image paths (same as A_paths)
"""
path = self.AB_paths[index]
im = Image.open(path).convert('RGB')
im = self.transform(im)
im = np.array(im)
lab = color.rgb2lab(im).astype(np.float32)
lab_t = transforms.ToTensor()(lab)
A = lab_t[[0], ...] / 50.0 - 1.0
B = lab_t[[1, 2], ...] / 110.0
return {'A': A, 'B': B, 'A_paths': path, 'B_paths': path}
def __len__(self):
"""Return the total number of images in the dataset."""
return len(self.AB_paths)

View File

@@ -0,0 +1,66 @@
"""A modified image folder class
We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
so that this class can load images from both current directory and its subdirectories.
"""
import torch.utils.data as data
from PIL import Image
import os
import os.path
IMG_EXTENSIONS = [
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def make_dataset(dir, max_dataset_size=float("inf")):
images = []
assert os.path.isdir(dir), '%s is not a valid directory' % dir
for root, _, fnames in sorted(os.walk(dir)):
for fname in fnames:
if is_image_file(fname):
path = os.path.join(root, fname)
images.append(path)
return images[:min(max_dataset_size, len(images))]
def default_loader(path):
return Image.open(path).convert('RGB')
class ImageFolder(data.Dataset):
def __init__(self, root, transform=None, return_paths=False,
loader=default_loader):
imgs = make_dataset(root)
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in: " + root + "\n"
"Supported image extensions are: " +
",".join(IMG_EXTENSIONS)))
self.root = root
self.imgs = imgs
self.transform = transform
self.return_paths = return_paths
self.loader = loader
def __getitem__(self, index):
path = self.imgs[index]
img = self.loader(path)
if self.transform is not None:
img = self.transform(img)
if self.return_paths:
return img, path
else:
return img
def __len__(self):
return len(self.imgs)

View File

@@ -0,0 +1,112 @@
import os
import ntpath
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
import random
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
class PairedDataset(BaseDataset):
"""
This dataset class can load paired datasets.
It requires two directories to host training images from domain A '/path/to/data/trainA'
and from domain B '/path/to/data/trainB' respectively.
You can train the model with the dataset flag '--dataroot /path/to/data'.
Similarly, you need to prepare two directories:
'/path/to/data/testA' and '/path/to/data/testB' during test time.
"""
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
parser.add_argument('--gt_prefix', type=str, default='', help='name of the used prior')
return parser
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A') # create a path '/path/to/data/trainA'
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B') # create a path '/path/to/data/trainB'
self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size)) # load images from '/path/to/data/trainA'
# self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size)) # load images from '/path/to/data/trainB'
self.A_size = len(self.A_paths) # get the size of dataset A
# self.B_size = len(self.B_paths) # get the size of dataset B
# btoA = self.opt.direction == 'BtoA'
# input_nc = self.opt.output_nc if btoA else self.opt.input_nc # get the number of channels of input image
# output_nc = self.opt.input_nc if btoA else self.opt.output_nc # get the number of channels of output image
self.transform = get_transform(self.opt, grayscale=(self.opt.input_nc == 1))
self.toTensor = transforms.ToTensor()
# self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index (int) -- a random integer for data indexing
Returns a dictionary that contains A, B, A_paths and B_paths
A (tensor) -- an image in the input domain
B (tensor) -- its corresponding image in the target domain
A_paths (str) -- image paths
B_paths (str) -- image paths
"""
A_path = self.A_paths[index] # make sure index is within then range
A_name = os.path.splitext(ntpath.basename(A_path))[0]
B_shortPath = '%s%s.png'%('_'.join(A_name.split('_')[:-1]), self.opt.gt_prefix) # '%s.png'%A_name.split('_')[0]
B_path = os.path.join(self.dir_B, B_shortPath)
A_img = Image.open(A_path).convert('RGB')
if os.path.exists(B_path):
B_img = Image.open(B_path).convert('RGB')
else:
print('file [%s] not exist!'%B_path)
B_img = A_img
if A_img.size != B_img.size:
B_img = self.cropImage(B_img, A_img.size)
# apply image transformation
A = self.transform(A_img)
# B = self.toTensor(B_img)
B = self.transform(B_img)
# return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path}
return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path,
'clear': B, 'paths': A_path}
def __len__(self):
"""Return the total number of images in the dataset.
As we have two datasets with potentially different number of images,
we take a maximum of
"""
return self.A_size
def cropImage(self, img, target_size):
ow, oh = img.size
tw, th = target_size
if (ow > tw or oh > th):
x1 = np.floor((ow - tw)/2)
y1 = np.floor((oh - th)/2)
return img.crop((x1, y1, x1 + tw, y1 + th))
return img

View File

@@ -0,0 +1,73 @@
### Copyright (C) 2017 NVIDIA Corporation. All rights reserved.
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
import os, ntpath
import numpy as np
from PIL import Image
import scipy.io as sio
import torchvision.transforms as transforms
from data.base_dataset import BaseDataset, get_params, get_transform
# from data.image_folder import make_dataset,
class SimpleBeDDEDataset(BaseDataset):
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
parser.add_argument('--bedde_list', required=True, type=str, help='image list of BeDDE')
return parser
def __init__(self, opt):
BaseDataset.__init__(self, opt)
self.data_list_file = opt.bedde_list
listFile = open(self.data_list_file, 'r')
self.imagePaths = listFile.read().split()
listFile.close()
self.I_size = len(self.imagePaths)
self.transform = get_transform(self.opt, grayscale=(self.opt.input_nc == 1))
self.toTensor = transforms.ToTensor()
def __getitem__(self, index):
### input A (label maps)
# print('bedde id %d'%index)
I_path = self.imagePaths[index]
I_img = Image.open(I_path).convert('RGB')
params = get_params(self.opt, I_img.size)
I_name = os.path.splitext(ntpath.basename(I_path))[0]
cityName = I_name.split('_')[0]
I_dir = ntpath.dirname(I_path)
base_dir = ntpath.dirname(I_dir)
J_path = os.path.join(base_dir, 'gt', '%s_clear.png'%cityName)
J_img = Image.open(J_path).convert('RGB')
base_dir = ntpath.dirname(I_dir)
mask_path = os.path.join(base_dir, 'mask', '%s_mask.mat'%I_name)
mask_info = sio.loadmat(mask_path)
J_root = ntpath.dirname(ntpath.dirname(I_path))
# apply image transformation
real_I = self.transform(I_img)
real_J = (self.toTensor(J_img) - 0.5) / 0.5
return {'haze': real_I, 'clear': real_J, 'mask': mask_info['mask'],
'city': cityName, 'paths': I_path}
# return {'haze': real_I , 'city': cityName, 'paths': curPath}
def __len__(self):
return self.I_size

View File

@@ -0,0 +1,110 @@
import os
import ntpath
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
import random
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
class SimplePairedDataset(BaseDataset):
"""
This dataset class can load paired datasets.
It requires two directories to host training images from domain A '/path/to/data/trainA'
and from domain B '/path/to/data/trainB' respectively.
You can train the model with the dataset flag '--dataroot /path/to/data'.
Similarly, you need to prepare two directories:
'/path/to/data/testA' and '/path/to/data/testB' during test time.
"""
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
parser.add_argument('--gt_prefix', type=str, default='', help='name of the used prior')
return parser
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A') # create a path '/path/to/data/trainA'
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B') # create a path '/path/to/data/trainB'
self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size)) # load images from '/path/to/data/trainA'
# self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size)) # load images from '/path/to/data/trainB'
self.A_size = len(self.A_paths) # get the size of dataset A
# self.B_size = len(self.B_paths) # get the size of dataset B
# btoA = self.opt.direction == 'BtoA'
# input_nc = self.opt.output_nc if btoA else self.opt.input_nc # get the number of channels of input image
# output_nc = self.opt.input_nc if btoA else self.opt.output_nc # get the number of channels of output image
self.transform = get_transform(self.opt, grayscale=(self.opt.input_nc == 1))
self.toTensor = transforms.ToTensor()
# self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index (int) -- a random integer for data indexing
Returns a dictionary that contains A, B, A_paths and B_paths
A (tensor) -- an image in the input domain
B (tensor) -- its corresponding image in the target domain
A_paths (str) -- image paths
B_paths (str) -- image paths
"""
A_path = self.A_paths[index] # make sure index is within then range
A_name = os.path.splitext(ntpath.basename(A_path))[0]
B_shortPath = '%s%s.png'%('_'.join(A_name.split('_')[:-1]), self.opt.gt_prefix) # '%s.png'%A_name.split('_')[0]
B_path = os.path.join(self.dir_B, B_shortPath)
A_img = Image.open(A_path).convert('RGB')
if os.path.exists(B_path):
B_img = Image.open(B_path).convert('RGB')
else:
print('file [%s] not exist!'%B_path)
B_img = A_img
if A_img.size != B_img.size:
B_img = self.cropImage(B_img, A_img.size)
# apply image transformation
A = self.transform(A_img)
B = (self.toTensor(B_img) - 0.5) / 0.5
# B = self.transform(B_img)
return {'haze': A, 'clear': B, 'paths': A_path, 'B_paths': B_path}
def __len__(self):
"""Return the total number of images in the dataset.
As we have two datasets with potentially different number of images,
we take a maximum of
"""
return self.A_size
def cropImage(self, img, target_size):
ow, oh = img.size
tw, th = target_size
if (ow > tw or oh > th):
x1 = np.floor((ow - tw)/2)
y1 = np.floor((oh - th)/2)
return img.crop((x1, y1, x1 + tw, y1 + th))
return img

View File

@@ -0,0 +1,45 @@
import os
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
class SingleDataset(BaseDataset):
"""This dataset class can load a set of images specified by the path --dataroot /path/to/data.
It can be used for generating CycleGAN results only for one side with the model option '-model test'.
"""
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase) # create a path '/path/to/data/testA'
self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))
input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
self.transform = get_transform(opt, grayscale=(input_nc == 1))
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index - - a random integer for data indexing
Returns a dictionary that contains A and A_paths
A(tensor) - - an image in one domain
A_paths(str) - - the path of the image
"""
A_path = self.A_paths[index]
A_img = Image.open(A_path).convert('RGB')
A = self.transform(A_img)
# return {'A': A, 'A_paths': A_path}
return {'haze': A, 'paths': A_path}
def __len__(self):
"""Return the total number of images in the dataset."""
return len(self.A_paths)

View File

@@ -0,0 +1,75 @@
"""Dataset class template
This module provides a template for users to implement custom datasets.
You can specify '--dataset_mode template' to use this dataset.
The class name should be consistent with both the filename and its dataset_mode option.
The filename should be <dataset_mode>_dataset.py
The class name should be <Dataset_mode>Dataset.py
You need to implement the following functions:
-- <modify_commandline_options>: Add dataset-specific options and rewrite default values for existing options.
-- <__init__>: Initialize this dataset class.
-- <__getitem__>: Return a data point and its metadata information.
-- <__len__>: Return the number of images.
"""
from data.base_dataset import BaseDataset, get_transform
# from data.image_folder import make_dataset
# from PIL import Image
class TemplateDataset(BaseDataset):
"""A template dataset class for you to implement custom datasets."""
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
parser.add_argument('--new_dataset_option', type=float, default=1.0, help='new dataset option')
parser.set_defaults(max_dataset_size=10, new_dataset_option=2.0) # specify dataset-specific default values
return parser
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
A few things can be done here.
- save the options (have been done in BaseDataset)
- get image paths and meta information of the dataset.
- define the image transformation.
"""
# save the option and dataset root
BaseDataset.__init__(self, opt)
# get the image paths of your dataset;
self.image_paths = [] # You can call sorted(make_dataset(self.root, opt.max_dataset_size)) to get all the image paths under the directory self.root
# define the default transform function. You can use <base_dataset.get_transform>; You can also define your custom transform function
self.transform = get_transform(opt)
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index -- a random integer for data indexing
Returns:
a dictionary of data with their names. It usually contains the data itself and its metadata information.
Step 1: get a random image path: e.g., path = self.image_paths[index]
Step 2: load your data from the disk: e.g., image = Image.open(path).convert('RGB').
Step 3: convert your data to a PyTorch tensor. You can use helpder functions such as self.transform. e.g., data = self.transform(image)
Step 4: return a data point as a dictionary.
"""
path = 'temp' # needs to be a string
data_A = None # needs to be a tensor
data_B = None # needs to be a tensor
return {'data_A': data_A, 'data_B': data_B, 'path': path}
def __len__(self):
"""Return the total number of images."""
return len(self.image_paths)

View File

@@ -0,0 +1,80 @@
import os.path
from data.base_dataset import BaseDataset, get_transform, get_params
from data.image_folder import make_dataset
from PIL import Image
import random
class UnalignedDataset(BaseDataset):
"""
This dataset class can load unaligned/unpaired datasets.
It requires two directories to host training images from domain A '/path/to/data/trainA'
and from domain B '/path/to/data/trainB' respectively.
You can train the model with the dataset flag '--dataroot /path/to/data'.
Similarly, you need to prepare two directories:
'/path/to/data/testA' and '/path/to/data/testB' during test time.
"""
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A') # create a path '/path/to/data/trainA'
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B') # create a path '/path/to/data/trainB'
self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size)) # load images from '/path/to/data/trainA'
self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size)) # load images from '/path/to/data/trainB'
self.A_size = len(self.A_paths) # get the size of dataset A
self.B_size = len(self.B_paths) # get the size of dataset B
# btoA = self.opt.direction == 'BtoA'
# input_nc = self.opt.output_nc if btoA else self.opt.input_nc # get the number of channels of input image
# output_nc = self.opt.input_nc if btoA else self.opt.output_nc # get the number of channels of output image
# self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))
# self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index (int) -- a random integer for data indexing
Returns a dictionary that contains A, B, A_paths and B_paths
A (tensor) -- an image in the input domain
B (tensor) -- its corresponding image in the target domain
A_paths (str) -- image paths
B_paths (str) -- image paths
"""
A_path = self.A_paths[index % self.A_size] # make sure index is within then range
if self.opt.serial_batches: # make sure index is within then range
index_B = index % self.B_size
else: # randomize the index for domain B to avoid fixed pairs.
index_B = random.randint(0, self.B_size - 1)
B_path = self.B_paths[index_B]
A_img = Image.open(A_path).convert('RGB')
B_img = Image.open(B_path).convert('RGB')
params_A = get_params(self.opt, A_img.size)
params_B = get_params(self.opt, B_img.size)
btoA = self.opt.direction == 'BtoA'
input_nc = self.opt.output_nc if btoA else self.opt.input_nc # get the number of channels of input image
output_nc = self.opt.input_nc if btoA else self.opt.output_nc # get the number of channels of output image
transform_A = get_transform(self.opt, params=params_A, grayscale=(input_nc == 1))
transform_B = get_transform(self.opt, params=params_B, grayscale=(output_nc == 1))
# apply image transformation
A = transform_A(A_img)
B = transform_B(B_img)
return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path}
def __len__(self):
"""Return the total number of images in the dataset.
As we have two datasets with potentially different number of images,
we take a maximum of
"""
return max(self.A_size, self.B_size)

View File

@@ -0,0 +1,72 @@
import os.path
from data.base_dataset import BaseDataset, get_transform, get_params
from data.image_folder import make_dataset
from PIL import Image
import random
class UnpairedDataset(BaseDataset):
"""
This dataset class can load unpaired datasets for dehazing.
It requires two directories to host training images from domain A '/path/to/data/trainA'
and from domain B '/path/to/data/trainB' respectively.
You can train the model with the dataset flag '--dataroot /path/to/data'.
Similarly, you need to prepare two directories:
'/path/to/data/testA' and '/path/to/data/testB' during test time.
"""
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.dir_I = os.path.join(opt.dataroot, opt.phase + 'A') # create a path '/path/to/data/trainA'
self.dir_J = os.path.join(opt.dataroot, opt.phase + 'B') # create a path '/path/to/data/trainB'
self.I_paths = sorted(make_dataset(self.dir_I, opt.max_dataset_size)) # load images from '/path/to/data/trainA'
self.J_paths = sorted(make_dataset(self.dir_J, opt.max_dataset_size)) # load images from '/path/to/data/trainB'
self.I_size = len(self.I_paths) # get the size of dataset A
self.J_size = len(self.J_paths) # get the size of dataset B
def __getitem__(self, index):
"""Return a data point and its metadata information.
Parameters:
index (int) -- a random integer for data indexing
Returns a dictionary that contains haze, clear, paths and J_paths
haze (tensor) -- hazy image
clear (tensor) -- clear image
paths (str) -- image paths
J_paths (str) -- image paths
"""
I_path = self.I_paths[index % self.I_size] # make sure index is within then range
if self.opt.serial_batches: # make sure index is within then range
index_J = index % self.J_size
else: # randomize the index for domain B to avoid fixed pairs.
index_J = random.randint(0, self.J_size - 1)
J_path = self.J_paths[index_J]
I_img = Image.open(I_path).convert('RGB')
J_img = Image.open(J_path).convert('RGB')
params_I = get_params(self.opt, I_img.size)
params_J = get_params(self.opt, J_img.size)
transform_I = get_transform(self.opt, params=params_I, grayscale=(self.opt.input_nc == 1))
transform_J = get_transform(self.opt, params=params_J, grayscale=(self.opt.output_nc == 1))
# apply image transformation
real_I = transform_I(I_img)
real_J = transform_J(J_img)
return {'haze': real_I, 'clear': real_J, 'paths': I_path, 'J_paths': J_path}
def __len__(self):
"""Return the total number of images in the dataset.
As we have two datasets with potentially different number of images,
we take a maximum of
"""
return max(self.I_size, self.J_size)

Binary file not shown.

After

Width:  |  Height:  |  Size: 196 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 196 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 689 KiB

View File

@@ -0,0 +1,67 @@
"""This package contains modules related to objective functions, optimizations, and network architectures.
To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
You need to implement the following five functions:
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
-- <set_input>: unpack data from dataset and apply preprocessing.
-- <forward>: produce intermediate results.
-- <optimize_parameters>: calculate loss, gradients, and update network weights.
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
In the function <__init__>, you need to define four lists:
-- self.loss_names (str list): specify the training losses that you want to plot and save.
-- self.model_names (str list): define networks used in our training.
-- self.visual_names (str list): specify the images that you want to display and save.
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
Now you can use the model class by specifying flag '--model dummy'.
See our template model class 'template_model.py' for more details.
"""
import importlib
from models.base_model import BaseModel
def find_model_using_name(model_name):
"""Import the module "models/[model_name]_model.py".
In the file, the class called DatasetNameModel() will
be instantiated. It has to be a subclass of BaseModel,
and it is case-insensitive.
"""
model_filename = "models." + model_name + "_model"
modellib = importlib.import_module(model_filename)
model = None
target_model_name = model_name.replace('_', '') + 'model'
for name, cls in modellib.__dict__.items():
if name.lower() == target_model_name.lower() \
and issubclass(cls, BaseModel):
model = cls
if model is None:
print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
exit(0)
return model
def get_option_setter(model_name):
"""Return the static method <modify_commandline_options> of the model class."""
model_class = find_model_using_name(model_name)
return model_class.modify_commandline_options
def create_model(opt):
"""Create a model given the option.
This function warps the class CustomDatasetDataLoader.
This is the main interface between this package and 'train.py'/'test.py'
Example:
>>> from models import create_model
>>> model = create_model(opt)
"""
model = find_model_using_name(opt.model)
instance = model(opt)
print("model [%s] was created" % type(instance).__name__)
return instance

View File

@@ -0,0 +1,229 @@
import os
import torch
from collections import OrderedDict
from abc import ABC, abstractmethod
from . import networks
class BaseModel(ABC):
"""This class is an abstract base class (ABC) for models.
To create a subclass, you need to implement the following five functions:
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
-- <set_input>: unpack data from dataset and apply preprocessing.
-- <forward>: produce intermediate results.
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
"""
def __init__(self, opt):
"""Initialize the BaseModel class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
When creating your custom class, you need to implement your own initialization.
In this fucntion, you should first call <BaseModel.__init__(self, opt)>
Then, you need to define four lists:
-- self.loss_names (str list): specify the training losses that you want to plot and save.
-- self.model_names (str list): specify the images that you want to display and save.
-- self.visual_names (str list): define networks used in our training.
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
"""
self.opt = opt
self.gpu_ids = opt.gpu_ids
self.isTrain = opt.isTrain
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # get device name: CPU or GPU
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) # save all the checkpoints to save_dir
if opt.preprocess != 'scale_width': # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
torch.backends.cudnn.benchmark = True
self.loss_names = []
self.model_names = []
self.visual_names = []
self.optimizers = []
self.image_paths = []
self.metric = 0 # used for learning rate policy 'plateau'
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new model-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
return parser
@abstractmethod
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): includes the data itself and its metadata information.
"""
pass
@abstractmethod
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
pass
@abstractmethod
def optimize_parameters(self):
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
pass
def setup(self, opt):
"""Load and print networks; create schedulers
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
if self.isTrain:
self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
if not self.isTrain or opt.continue_train:
load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
self.load_networks(load_suffix)
self.print_networks(opt.verbose)
def eval(self):
"""Make models eval mode during test time"""
for name in self.model_names:
if isinstance(name, str):
net = getattr(self, 'net' + name)
net.eval()
def test(self):
"""Forward function used in test time.
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
It also calls <compute_visuals> to produce additional visualization results
"""
with torch.no_grad():
self.forward()
self.compute_visuals()
def compute_visuals(self):
"""Calculate additional output images for visdom and HTML visualization"""
pass
def get_image_paths(self):
""" Return image paths that are used to load current data"""
return self.image_paths
def update_learning_rate(self):
"""Update learning rates for all the networks; called at the end of every epoch"""
for scheduler in self.schedulers:
if self.opt.lr_policy == 'plateau':
scheduler.step(self.metric)
else:
scheduler.step()
lr = self.optimizers[0].param_groups[0]['lr']
print('learning rate = %.7f' % lr)
def get_current_visuals(self):
"""Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
visual_ret = OrderedDict()
for name in self.visual_names:
if isinstance(name, str):
visual_ret[name] = getattr(self, name)
return visual_ret
def get_current_losses(self):
"""Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
errors_ret = OrderedDict()
for name in self.loss_names:
if isinstance(name, str):
errors_ret[name] = float(getattr(self, 'loss_' + name)) # float(...) works for both scalar tensor and float number
return errors_ret
def save_networks(self, epoch):
"""Save all the networks to the disk.
Parameters:
epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
"""
for name in self.model_names:
if isinstance(name, str):
save_filename = '%s_net_%s.pth' % (epoch, name)
save_path = os.path.join(self.save_dir, save_filename)
net = getattr(self, 'net' + name)
if len(self.gpu_ids) > 0 and torch.cuda.is_available():
torch.save(net.module.cpu().state_dict(), save_path)
net.cuda(self.gpu_ids[0])
else:
torch.save(net.cpu().state_dict(), save_path)
def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
"""Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
key = keys[i]
if i + 1 == len(keys): # at the end, pointing to a parameter/buffer
if module.__class__.__name__.startswith('InstanceNorm') and \
(key == 'running_mean' or key == 'running_var'):
if getattr(module, key) is None:
state_dict.pop('.'.join(keys))
if module.__class__.__name__.startswith('InstanceNorm') and \
(key == 'num_batches_tracked'):
state_dict.pop('.'.join(keys))
else:
self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
def load_networks(self, epoch):
"""Load all the networks from the disk.
Parameters:
epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
"""
for name in self.model_names:
if isinstance(name, str):
load_filename = '%s_net_%s.pth' % (epoch, name)
load_path = os.path.join(self.save_dir, load_filename)
net = getattr(self, 'net' + name)
if isinstance(net, torch.nn.DataParallel):
net = net.module
print('loading the model from %s' % load_path)
# if you are using PyTorch newer than 0.4 (e.g., built from
# GitHub source), you can remove str() on self.device
state_dict = torch.load(load_path, map_location=str(self.device))
if hasattr(state_dict, '_metadata'):
del state_dict._metadata
# patch InstanceNorm checkpoints prior to 0.4
for key in list(state_dict.keys()): # need to copy keys here because we mutate in loop
self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
net.load_state_dict(state_dict)
def print_networks(self, verbose):
"""Print the total number of parameters in the network and (if verbose) network architecture
Parameters:
verbose (bool) -- if verbose: print the network architecture
"""
print('---------- Networks initialized -------------')
for name in self.model_names:
if isinstance(name, str):
net = getattr(self, 'net' + name)
num_params = 0
for param in net.parameters():
num_params += param.numel()
if verbose:
print(net)
print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
print('-----------------------------------------------')
def set_requires_grad(self, nets, requires_grad=False):
"""Set requies_grad=Fasle for all the networks to avoid unnecessary computations
Parameters:
nets (network list) -- a list of networks
requires_grad (bool) -- whether the networks require gradients or not
"""
if not isinstance(nets, list):
nets = [nets]
for net in nets:
if net is not None:
for param in net.parameters():
param.requires_grad = requires_grad

View File

@@ -0,0 +1,221 @@
import torch
import itertools
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks
import torch.nn.functional as F
from util import util
class BasicDehazeModel(BaseModel):
"""
This class implements the CycleGAN model, for learning image-to-image translation without paired data.
The model training requires '--dataset_mode unaligned' dataset.
By default, it uses a '--netG resnet_9blocks' ResNet generator,
a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
and a least-square GANs objective ('--gan_mode lsgan').
CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
A (source domain), B (target domain).
Generators: G_A: A -> B; G_B: B -> A.
Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
Forward cycle loss: lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
Dropout is not used in the original CycleGAN paper.
"""
parser.set_defaults(no_dropout=True) # default CycleGAN did not use dropout
if is_train:
parser.add_argument('--lambda_haze', type=float, default=0.1, help='weight for D_haze')
parser.add_argument('--lambda_clear', type=float, default=0.1, help='weight for D_clear')
parser.add_argument('--lambda_tv', type=float, default=1, help='weight for D_clear')
parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
parser.add_argument('--netR_T', type=str, default='unet_trans_256', help='specify generator architecture')
parser.add_argument('--netR_J', type=str, default='haze_refine_2', help='specify generator architecture')
return parser
def __init__(self, opt):
"""Initialize the CycleGAN class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['D_haze', 'G_rec_I', 'D_clear', 'G_ref_J', 'rec_I', 'rec_J', 'TV_T', 'idt_J']
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['real_I', 'est_J', 'rec_I', 'rec_J',
'est_T_vis', 'out_T_vis', 'real_J']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
if self.isTrain:
self.model_names = ['Est_T', 'Est_J', 'D_haze', 'D_clear']
else: # during test time, only load Gs
self.model_names = ['Est_T', 'Est_J']
# define networks (both Generators and discriminators)
self.netEst_T = networks.define_G(opt.input_nc, 1, opt.ngf, opt.netR_T, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
self.netEst_J = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netR_J, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain: # define discriminators
self.netD_haze = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
self.netD_clear = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
if opt.lambda_identity > 0.0: # only works when input and output images have the same number of channels
assert(opt.input_nc == opt.output_nc)
self.fake_I_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
self.fake_J_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
# # define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) # define GAN loss.
self.criterionRec = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.criterionTV = networks.TVLoss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(itertools.chain(self.netEst_T.parameters(), self.netEst_J.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_haze.parameters(), self.netD_clear.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): include the data itself and its metadata information.
The option 'direction' can be used to swap domain A and domain B.
"""
self.real_I = input['haze'].to(self.device) # [-1, 1]
self.real_J = input['clear'].to(self.device) # [-1, 1]
self.image_paths = input['paths']
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
# output scale [0,1]
self.est_T, self.out_T = self.netEst_T(self.real_I)
self.est_J = self.netEst_J(self.real_I)
# reconstruct haze image
est_T_map = self.est_T.repeat(1,3,1,1)
self.rec_I = util.synthesize_fog(self.est_J, est_T_map)
self.rec_J = util.reverse_fog(self.real_I, est_T_map)
def test(self):
"""Forward function used in test time.
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
It also calls <compute_visuals> to produce additional visualization results
"""
with torch.no_grad():
self.forward()
self.compute_visuals()
self.refine_J = (self.rec_J + self.est_J)/2
self.visual_names.append('refine_J')
def compute_visuals(self):
"""Calculate additional output images for visdom and HTML visualization"""
# rescale to [-1,1] for visdom
self.est_T_vis = (self.est_T - 0.5)/0.5
self.out_T_vis = (self.out_T - 0.5)/0.5
# self.map_A_vis = (self.map_A - 0.5)/0.5
def backward_D_basic(self, netD, real, fake):
"""Calculate GAN loss for the discriminator
Parameters:
netD (network) -- the discriminator D
real (tensor array) -- real images
fake (tensor array) -- images generated by a generator
Return the discriminator loss.
We also call loss_D.backward() to calculate the gradients.
"""
# Real
pred_real = netD(real)
loss_D_real = self.criterionGAN(pred_real, True)
# Fake
pred_fake = netD(fake.detach())
loss_D_fake = self.criterionGAN(pred_fake, False)
# Combined loss and calculate gradients
loss_D = (loss_D_real + loss_D_fake) * 0.5
loss_D.backward()
return loss_D
def backward_D_haze(self):
fake_I = self.fake_I_pool.query(self.rec_I)
self.loss_D_haze = self.backward_D_basic(self.netD_haze, self.real_I, fake_I)
def backward_D_clear(self):
fake_J = self.fake_J_pool.query(self.est_J)
self.loss_D_clear = self.backward_D_basic(self.netD_clear, self.real_J, fake_J)
def backward_G(self):
lambda_idt = self.opt.lambda_identity
lambda_tv = self.opt.lambda_tv
lambda_haze = self.opt.lambda_haze
lambda_clear = self.opt.lambda_clear
# TV loss
if lambda_tv > 0.0:
self.loss_TV_T = self.criterionTV(self.out_T)*lambda_tv
else:
self.loss_TV_T = 0
# Identity loss
if lambda_idt > 0.0:
self.loss_idt_J = self.criterionIdt(self.netEst_J(self.real_J), self.real_J)*lambda_idt
else:
self.loss_idt_J = 0
# Generator losses for rec_I and est_J
self.loss_G_rec_I = self.criterionGAN(self.netD_haze(self.rec_I), True)*lambda_haze
self.loss_G_ref_J = self.criterionGAN(self.netD_clear(self.est_J), True)*lambda_clear #+ \
# self.criterionGAN(self.netD_clear(self.rec_J), True)*lambda_clear
# Reconstrcut loss
self.loss_rec_I = self.criterionRec(self.rec_I, self.real_I)
# only compute, not back propagate
self.loss_rec_J = self.criterionRec(self.rec_J, self.est_J)
self.loss_G = self.loss_G_rec_I + self.loss_G_ref_J + self.loss_rec_I + self.loss_idt_J + self.loss_TV_T
self.loss_G.backward()
def optimize_parameters(self):
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
# forward
self.forward() # compute fake images and reconstruction images.
# G_A and G_B
self.set_requires_grad([self.netD_haze, self.netD_clear], False) # Ds require no gradients when optimizing Gs
self.optimizer_G.zero_grad() # set G_A and G_B's gradients to zero
self.backward_G() # calculate gradients for G_A and G_B
self.optimizer_G.step() # update G_A and G_B's weights
# D_A and D_B
self.set_requires_grad([self.netD_haze, self.netD_clear], True)
self.optimizer_D.zero_grad() # set D_A and D_B's gradients to zero
self.backward_D_haze() # calculate gradients for D_A
self.backward_D_clear() # calculate graidents for D_B
self.optimizer_D.step() # update D_A and D_B's weights

View File

@@ -0,0 +1,68 @@
from .pix2pix_model import Pix2PixModel
import torch
from skimage import color # used for lab2rgb
import numpy as np
class ColorizationModel(Pix2PixModel):
"""This is a subclass of Pix2PixModel for image colorization (black & white image -> colorful images).
The model training requires '-dataset_model colorization' dataset.
It trains a pix2pix model, mapping from L channel to ab channels in Lab color space.
By default, the colorization dataset will automatically set '--input_nc 1' and '--output_nc 2'.
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
By default, we use 'colorization' dataset for this model.
See the original pix2pix paper (https://arxiv.org/pdf/1611.07004.pdf) and colorization results (Figure 9 in the paper)
"""
Pix2PixModel.modify_commandline_options(parser, is_train)
parser.set_defaults(dataset_mode='colorization')
return parser
def __init__(self, opt):
"""Initialize the class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
For visualization, we set 'visual_names' as 'real_A' (input real image),
'real_B_rgb' (ground truth RGB image), and 'fake_B_rgb' (predicted RGB image)
We convert the Lab image 'real_B' (inherited from Pix2pixModel) to a RGB image 'real_B_rgb'.
we convert the Lab image 'fake_B' (inherited from Pix2pixModel) to a RGB image 'fake_B_rgb'.
"""
# reuse the pix2pix model
Pix2PixModel.__init__(self, opt)
# specify the images to be visualized.
self.visual_names = ['real_A', 'real_B_rgb', 'fake_B_rgb']
def lab2rgb(self, L, AB):
"""Convert an Lab tensor image to a RGB numpy output
Parameters:
L (1-channel tensor array): L channel images (range: [-1, 1], torch tensor array)
AB (2-channel tensor array): ab channel images (range: [-1, 1], torch tensor array)
Returns:
rgb (RGB numpy image): rgb output images (range: [0, 255], numpy array)
"""
AB2 = AB * 110.0
L2 = (L + 1.0) * 50.0
Lab = torch.cat([L2, AB2], dim=1)
Lab = Lab[0].data.cpu().float().numpy()
Lab = np.transpose(Lab.astype(np.float64), (1, 2, 0))
rgb = color.lab2rgb(Lab) * 255
return rgb
def compute_visuals(self):
"""Calculate additional output images for visdom and HTML visualization"""
self.real_B_rgb = self.lab2rgb(self.real_A, self.real_B)
self.fake_B_rgb = self.lab2rgb(self.real_A, self.fake_B)

View File

@@ -0,0 +1,211 @@
import torch
import itertools
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks
class CycleGANModel(BaseModel):
"""
This class implements the CycleGAN model, for learning image-to-image translation without paired data.
The model training requires '--dataset_mode unaligned' dataset.
By default, it uses a '--netG resnet_9blocks' ResNet generator,
a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
and a least-square GANs objective ('--gan_mode lsgan').
CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
A (source domain), B (target domain).
Generators: G_A: A -> B; G_B: B -> A.
Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
Forward cycle loss: lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
Dropout is not used in the original CycleGAN paper.
"""
parser.set_defaults(no_dropout=True) # default CycleGAN did not use dropout
if is_train:
parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
return parser
def __init__(self, opt):
"""Initialize the CycleGAN class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
visual_names_A = ['real_A', 'fake_B', 'rec_A']
visual_names_B = ['real_B', 'fake_A', 'rec_B']
if self.isTrain and self.opt.lambda_identity > 0.0: # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B)
visual_names_A.append('idt_B')
visual_names_B.append('idt_A')
self.visual_names = visual_names_A + visual_names_B # combine visualizations for A and B
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
if self.isTrain:
self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
else: # during test time, only load Gs
self.model_names = ['G_A', 'G_B']
# define networks (both Generators and discriminators)
# The naming is different from those used in the paper.
# Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain: # define discriminators
self.netD_A = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
if opt.lambda_identity > 0.0: # only works when input and output images have the same number of channels
assert(opt.input_nc == opt.output_nc)
self.fake_A_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
self.fake_B_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) # define GAN loss.
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): include the data itself and its metadata information.
The option 'direction' can be used to swap domain A and domain B.
"""
if hasattr(self.opt, 'prior_name'):
self.real_A = input['haze'].to(self.device)
self.real_B = input['clear'].to(self.device)
self.image_paths = input['paths']
else:
AtoB = self.opt.direction == 'AtoB'
self.real_A = input['A' if AtoB else 'B'].to(self.device)
self.real_B = input['B' if AtoB else 'A'].to(self.device)
self.image_paths = input['A_paths' if AtoB else 'B_paths']
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
self.fake_B = self.netG_A(self.real_A) # G_A(A)
self.rec_A = self.netG_B(self.fake_B) # G_B(G_A(A))
self.fake_A = self.netG_B(self.real_B) # G_B(B)
self.rec_B = self.netG_A(self.fake_A) # G_A(G_B(B))
def test(self):
"""Forward function used in test time.
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
It also calls <compute_visuals> to produce additional visualization results
"""
with torch.no_grad():
self.forward()
self.compute_visuals()
self.visual_names.append('refine_J')
self.refine_J = self.fake_B
def backward_D_basic(self, netD, real, fake):
"""Calculate GAN loss for the discriminator
Parameters:
netD (network) -- the discriminator D
real (tensor array) -- real images
fake (tensor array) -- images generated by a generator
Return the discriminator loss.
We also call loss_D.backward() to calculate the gradients.
"""
# Real
pred_real = netD(real)
loss_D_real = self.criterionGAN(pred_real, True)
# Fake
pred_fake = netD(fake.detach())
loss_D_fake = self.criterionGAN(pred_fake, False)
# Combined loss and calculate gradients
loss_D = (loss_D_real + loss_D_fake) * 0.5
loss_D.backward()
return loss_D
def backward_D_A(self):
"""Calculate GAN loss for discriminator D_A"""
fake_B = self.fake_B_pool.query(self.fake_B)
self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)
def backward_D_B(self):
"""Calculate GAN loss for discriminator D_B"""
fake_A = self.fake_A_pool.query(self.fake_A)
self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)
def backward_G(self):
"""Calculate the loss for generators G_A and G_B"""
lambda_idt = self.opt.lambda_identity
lambda_A = self.opt.lambda_A
lambda_B = self.opt.lambda_B
# Identity loss
if lambda_idt > 0:
# G_A should be identity if real_B is fed: ||G_A(B) - B||
self.idt_A = self.netG_A(self.real_B)
self.loss_idt_A = self.criterionIdt(self.idt_A, self.real_B) * lambda_B * lambda_idt
# G_B should be identity if real_A is fed: ||G_B(A) - A||
self.idt_B = self.netG_B(self.real_A)
self.loss_idt_B = self.criterionIdt(self.idt_B, self.real_A) * lambda_A * lambda_idt
else:
self.loss_idt_A = 0
self.loss_idt_B = 0
# GAN loss D_A(G_A(A))
self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)
# GAN loss D_B(G_B(B))
self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)
# Forward cycle loss || G_B(G_A(A)) - A||
self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A
# Backward cycle loss || G_A(G_B(B)) - B||
self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B
# combined loss and calculate gradients
self.loss_G = self.loss_G_A + self.loss_G_B + self.loss_cycle_A + self.loss_cycle_B + self.loss_idt_A + self.loss_idt_B
self.loss_G.backward()
def optimize_parameters(self):
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
# forward
self.forward() # compute fake images and reconstruction images.
# G_A and G_B
self.set_requires_grad([self.netD_A, self.netD_B], False) # Ds require no gradients when optimizing Gs
self.optimizer_G.zero_grad() # set G_A and G_B's gradients to zero
self.backward_G() # calculate gradients for G_A and G_B
self.optimizer_G.step() # update G_A and G_B's weights
# D_A and D_B
self.set_requires_grad([self.netD_A, self.netD_B], True)
self.optimizer_D.zero_grad() # set D_A and D_B's gradients to zero
self.backward_D_A() # calculate gradients for D_A
self.backward_D_B() # calculate graidents for D_B
self.optimizer_D.step() # update D_A and D_B's weights

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,127 @@
import torch
from .base_model import BaseModel
from . import networks
class Pix2PixModel(BaseModel):
""" This class implements the pix2pix model, for learning a mapping from input images to output images given paired data.
The model training requires '--dataset_mode aligned' dataset.
By default, it uses a '--netG unet256' U-Net generator,
a '--netD basic' discriminator (PatchGAN),
and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper).
pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
For pix2pix, we do not use image buffer
The training objective is: GAN Loss + lambda_L1 * ||G(A)-B||_1
By default, we use vanilla GAN loss, UNet with batchnorm, and aligned datasets.
"""
# changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
parser.set_defaults(norm='batch', netG='unet_256', dataset_mode='aligned')
if is_train:
parser.set_defaults(pool_size=0, gan_mode='vanilla')
parser.add_argument('--lambda_L1', type=float, default=100.0, help='weight for L1 loss')
return parser
def __init__(self, opt):
"""Initialize the pix2pix class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['real_A', 'fake_B', 'real_B']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load G
self.model_names = ['G']
# define networks (both generator and discriminator)
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
self.criterionL1 = torch.nn.L1Loss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): include the data itself and its metadata information.
The option 'direction' can be used to swap images in domain A and domain B.
"""
AtoB = self.opt.direction == 'AtoB'
self.real_A = input['A' if AtoB else 'B'].to(self.device)
self.real_B = input['B' if AtoB else 'A'].to(self.device)
self.image_paths = input['A_paths' if AtoB else 'B_paths']
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
self.fake_B = self.netG(self.real_A) # G(A)
def backward_D(self):
"""Calculate GAN loss for the discriminator"""
# Fake; stop backprop to the generator by detaching fake_B
fake_AB = torch.cat((self.real_A, self.fake_B), 1) # we use conditional GANs; we need to feed both input and output to the discriminator
pred_fake = self.netD(fake_AB.detach())
self.loss_D_fake = self.criterionGAN(pred_fake, False)
# Real
real_AB = torch.cat((self.real_A, self.real_B), 1)
pred_real = self.netD(real_AB)
self.loss_D_real = self.criterionGAN(pred_real, True)
# combine loss and calculate gradients
self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
self.loss_D.backward()
def backward_G(self):
"""Calculate GAN and L1 loss for the generator"""
# First, G(A) should fake the discriminator
fake_AB = torch.cat((self.real_A, self.fake_B), 1)
pred_fake = self.netD(fake_AB)
self.loss_G_GAN = self.criterionGAN(pred_fake, True)
# Second, G(A) = B
self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
# combine loss and calculate gradients
self.loss_G = self.loss_G_GAN + self.loss_G_L1
self.loss_G.backward()
def optimize_parameters(self):
self.forward() # compute fake images: G(A)
# update D
self.set_requires_grad(self.netD, True) # enable backprop for D
self.optimizer_D.zero_grad() # set D's gradients to zero
self.backward_D() # calculate gradients for D
self.optimizer_D.step() # update D's weights
# update G
self.set_requires_grad(self.netD, False) # D requires no gradients when optimizing G
self.optimizer_G.zero_grad() # set G's gradients to zero
self.backward_G() # calculate graidents for G
self.optimizer_G.step() # udpate G's weights

View File

@@ -0,0 +1,226 @@
import torch
import itertools
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks
import torch.nn.functional as F
from util import util
class RefinedDCPModel(BaseModel):
"""
This class implements the RefineDNet model, for learning single image dehazing without paired data.
It adopts the basic backbone networks provided by CycleGAN.
The model training requires '--dataset_mode unpaired' dataset.
By default, it uses a '--netR_T unet_trans_256' U-Net refiner,
a '--netR_J resnet_9blocks' ResNet refiner,
and a '--netD basic' discriminator (PatchGAN introduced by pix2pix).
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
parser.set_defaults(no_dropout=True) # default CycleGAN did not use dropout
if is_train:
parser.add_argument('--lambda_G', type=float, default=0.05, help='weight for loss_G_single')
parser.add_argument('--lambda_identity', type=float, default=1, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
parser.add_argument('--lambda_rec_I', type=float, default=1, help='weight for loss_rec_I')
parser.add_argument('--lambda_tv', type=float, default=1, help='weight for TV loss of refine_T')
parser.add_argument('--lambda_vgg', type=float, default=0, help='weight for loss_vgg')
parser.add_argument('--netR_T', type=str, default='unet_trans_256', help='specify generator architecture')
parser.add_argument('--netR_J', type=str, default='resnet_9blocks', help='specify generator architecture')
return parser
def __init__(self, opt):
"""Initialize the RefineDNet class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['D_single', 'G_single', 'rec_I', 'TV_T', 'idt_J', 'vgg']
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
if self.isTrain:
self.visual_names = ['real_I', 'dcp_T_vis', 'refine_T_vis', 'out_T_vis', 'dcp_J','refine_J', 'rec_I', 'rec_J','map_A', 'real_J', 'ref_real_J']
else:
self.visual_names = ['real_I', 'dcp_T_vis', 'refine_T_vis', 'out_T_vis', 'dcp_J','refine_J', 'rec_I', 'rec_J','map_A']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
if self.isTrain:
self.model_names = ['Refiner_T', 'Refiner_J', 'D']
else: # during test time, only load Gs
self.model_names = ['Refiner_T', 'Refiner_J']
# define networks (both Generators and discriminators)
self.netG_DCP = networks.init_net(networks.DCPDehazeGenerator(), gpu_ids=self.gpu_ids) # use default setting for DCP
self.netRefiner_T = networks.define_G(opt.input_nc+1, 1, opt.ngf, opt.netR_T, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
self.netRefiner_J = networks.define_G(opt.input_nc+opt.output_nc, opt.output_nc, opt.ngf, opt.netR_J, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain: # define discriminators
self.netD = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
if opt.lambda_identity > 0.0: # only works when input and output images have the same number of channels
assert(opt.input_nc == opt.output_nc)
self.fake_I_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
self.fake_J_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
# # define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) # define GAN loss.
self.criterionRec = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.criterionTV = networks.TVLoss()
self.criterionVGG = networks.VGGLoss() if self.opt.lambda_vgg > 0.0 else None
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(itertools.chain(self.netRefiner_T.parameters(), self.netRefiner_J.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
# display the architecture of each part
# print(self.netRefiner_T)
# print(self.netRefiner_J)
# if self.isTrain:
# print(self.netD)
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): include the data itself and its metadata information.
"""
self.real_I = input['haze'].to(self.device) # [-1, 1]
self.image_paths = input['paths']
if self.isTrain:
self.real_J = input['clear'].to(self.device) # [-1, 1]
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
dcp_J, self.dcp_T, self.dcp_A = self.netG_DCP(self.real_I)
#scale to [-1,1]
self.dcp_J = (torch.clamp(dcp_J,0,1)-0.5)/0.5
# output scale [0,1]
self.refine_T, self.out_T = self.netRefiner_T(torch.cat((self.real_I, self.dcp_T), 1))
self.refine_J = self.netRefiner_J(torch.cat((self.real_I, self.dcp_J), 1))
# reconstruct haze image
shape = self.refine_J.shape
dcp_A_scale = self.dcp_A
self.map_A = (dcp_A_scale).reshape((1,3,1,1)).repeat(1,1,shape[2],shape[3])
refine_T_map = self.refine_T.repeat(1,3,1,1)
self.rec_I = util.synthesize_fog(self.refine_J, refine_T_map, self.map_A)
self.rec_J = util.reverse_fog(self.real_I, refine_T_map, self.map_A)
def test(self):
"""Forward function used in test time.
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
It also calls <compute_visuals> to produce additional visualization results
"""
with torch.no_grad():
self.forward()
self.compute_visuals()
def compute_visuals(self):
"""Calculate additional output images for visdom and HTML visualization"""
# rescale to [-1,1] for visdom
self.refine_T_vis = (self.refine_T - 0.5)/0.5
self.out_T_vis = (self.out_T - 0.5)/0.5
self.dcp_T_vis = (self.dcp_T - 0.5)/0.5
# self.map_A_vis = (self.map_A - 0.5)/0.5
def backward_D_basic(self, netD, real, fake):
"""Calculate GAN loss for the discriminator
Parameters:
netD (network) -- the discriminator D
real (tensor array) -- real images
fake (tensor array) -- images generated by a generator
Return the discriminator loss.
We also call loss_D.backward() to calculate the gradients.
"""
# Real
pred_real = netD(real)
loss_D_real = self.criterionGAN(pred_real, True)
# Fake
pred_fake = netD(fake.detach())
loss_D_fake = self.criterionGAN(pred_fake, False)
# Combined loss and calculate gradients
loss_D = (loss_D_real + loss_D_fake) * 0.5
loss_D.backward()
return loss_D
def backward_D(self):
fake_J = self.fake_I_pool.query(self.refine_J)
self.loss_D_single = self.backward_D_basic(self.netD, self.real_J, fake_J)
def backward_G(self):
lambda_idt = self.opt.lambda_identity
lambda_tv = self.opt.lambda_tv
lambda_G = self.opt.lambda_G
lambda_rec_I = self.opt.lambda_rec_I
lambda_vgg = self.opt.lambda_vgg
# Generator losses for rec_I and refine_J
self.loss_G_single = self.criterionGAN(self.netD(self.refine_J), True)*lambda_G
# Reconstrcut loss
self.loss_rec_I = self.criterionRec(self.rec_I, self.real_I) * lambda_rec_I
# perecptual loss
self.loss_vgg = self.criterionVGG(self.refine_J, self.dcp_J)*lambda_vgg if lambda_vgg > 0.0 else 0
# TV loss
self.loss_TV_T = self.criterionTV(self.out_T)*lambda_tv if lambda_tv > 0.0 else 0
# Identity loss, ||refiner_J(real_J) - real_J||
self.ref_real_J = self.netRefiner_J(torch.cat((self.real_I, self.real_J), 1))
self.loss_idt_J = self.criterionIdt(self.ref_real_J, self.real_J)*lambda_idt \
if lambda_idt > 0.0 \
else 0
self.loss_G = self.loss_G_single + self.loss_rec_I + self.loss_idt_J \
+ self.loss_TV_T \
+ self.loss_vgg
self.loss_G.backward()
def optimize_parameters(self):
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
# forward
self.forward() # compute fake images and reconstruction images.
# G_A and G_B
self.set_requires_grad(self.netD, False) # Ds require no gradients when optimizing Gs
self.optimizer_G.zero_grad() # set G_A and G_B's gradients to zero
self.backward_G() # calculate gradients for G_A and G_B
self.optimizer_G.step() # update G_A and G_B's weights
# D_A and D_B
self.set_requires_grad(self.netD, True)
self.optimizer_D.zero_grad() # set D_A and D_B's gradients to zero
self.backward_D() # calculate gradients for D_A
self.optimizer_D.step() # update D_A and D_B's weights

View File

@@ -0,0 +1,99 @@
"""Model class template
This module provides a template for users to implement custom models.
You can specify '--model template' to use this model.
The class name should be consistent with both the filename and its model option.
The filename should be <model>_dataset.py
The class name should be <Model>Dataset.py
It implements a simple image-to-image translation baseline based on regression loss.
Given input-output pairs (data_A, data_B), it learns a network netG that can minimize the following L1 loss:
min_<netG> ||netG(data_A) - data_B||_1
You need to implement the following functions:
<modify_commandline_options>: Add model-specific options and rewrite default values for existing options.
<__init__>: Initialize this model class.
<set_input>: Unpack input data and perform data pre-processing.
<forward>: Run forward pass. This will be called by both <optimize_parameters> and <test>.
<optimize_parameters>: Update network weights; it will be called in every training iteration.
"""
import torch
from .base_model import BaseModel
from . import networks
class TemplateModel(BaseModel):
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new model-specific options and rewrite default values for existing options.
Parameters:
parser -- the option parser
is_train -- if it is training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
parser.set_defaults(dataset_mode='aligned') # You can rewrite default values for this model. For example, this model usually uses aligned dataset as its dataset.
if is_train:
parser.add_argument('--lambda_regression', type=float, default=1.0, help='weight for the regression loss') # You can define new arguments for this model.
return parser
def __init__(self, opt):
"""Initialize this model class.
Parameters:
opt -- training/test options
A few things can be done here.
- (required) call the initialization function of BaseModel
- define loss function, visualization images, model names, and optimizers
"""
BaseModel.__init__(self, opt) # call the initialization method of BaseModel
# specify the training losses you want to print out. The program will call base_model.get_current_losses to plot the losses to the console and save them to the disk.
self.loss_names = ['loss_G']
# specify the images you want to save and display. The program will call base_model.get_current_visuals to save and display these images.
self.visual_names = ['data_A', 'data_B', 'output']
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks to save and load networks.
# you can use opt.isTrain to specify different behaviors for training and test. For example, some networks will not be used during test, and you don't need to load them.
self.model_names = ['G']
# define networks; you can use opt.isTrain to specify different behaviors for training and test.
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, gpu_ids=self.gpu_ids)
if self.isTrain: # only defined during training time
# define your loss functions. You can use losses provided by torch.nn such as torch.nn.L1Loss.
# We also provide a GANLoss class "networks.GANLoss". self.criterionGAN = networks.GANLoss().to(self.device)
self.criterionLoss = torch.nn.L1Loss()
# define and initialize optimizers. You can define one optimizer for each network.
# If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
self.optimizer = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers = [self.optimizer]
# Our program will automatically call <model.setup> to define schedulers, load networks, and print networks
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input: a dictionary that contains the data itself and its metadata information.
"""
AtoB = self.opt.direction == 'AtoB' # use <direction> to swap data_A and data_B
self.data_A = input['A' if AtoB else 'B'].to(self.device) # get image data A
self.data_B = input['B' if AtoB else 'A'].to(self.device) # get image data B
self.image_paths = input['A_paths' if AtoB else 'B_paths'] # get image paths
def forward(self):
"""Run forward pass. This will be called by both functions <optimize_parameters> and <test>."""
self.output = self.netG(self.data_A) # generate output image given the input data_A
def backward(self):
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
# caculate the intermediate results if necessary; here self.output has been computed during function <forward>
# calculate loss given the input and intermediate results
self.loss_G = self.criterionLoss(self.output, self.data_B) * self.opt.lambda_regression
self.loss_G.backward() # calculate gradients of network G w.r.t. loss_G
def optimize_parameters(self):
"""Update network weights; it will be called in every training iteration."""
self.forward() # first call forward to calculate intermediate results
self.optimizer.zero_grad() # clear network G's existing gradients
self.backward() # calculate gradients for network G
self.optimizer.step() # update gradients for network G

View File

@@ -0,0 +1,69 @@
from .base_model import BaseModel
from . import networks
class TestModel(BaseModel):
""" This TesteModel can be used to generate CycleGAN results for only one direction.
This model will automatically set '--dataset_mode single', which only loads the images from one collection.
See the test instruction for more details.
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
The model can only be used during test time. It requires '--dataset_mode single'.
You need to specify the network using the option '--model_suffix'.
"""
assert not is_train, 'TestModel cannot be used during training time'
parser.set_defaults(dataset_mode='single')
parser.add_argument('--model_suffix', type=str, default='', help='In checkpoints_dir, [epoch]_net_G[model_suffix].pth will be loaded as the generator.')
return parser
def __init__(self, opt):
"""Initialize the pix2pix class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
assert(not opt.isTrain)
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = []
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['real', 'fake']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
self.model_names = ['G' + opt.model_suffix] # only generator is needed.
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG,
opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
# assigns the model to self.netG_[suffix] so that it can be loaded
# please see <BaseModel.load_networks>
setattr(self, 'netG' + opt.model_suffix, self.netG) # store netG in self.
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input: a dictionary that contains the data itself and its metadata information.
We need to use 'single_dataset' dataset mode. It only load images from one domain.
"""
self.real = input['A'].to(self.device)
self.image_paths = input['A_paths']
def forward(self):
"""Run forward pass."""
self.fake = self.netG(self.real) # G(real)
def optimize_parameters(self):
"""No optimization for test model."""
pass

View File

@@ -0,0 +1 @@
"""This package options includes option modules: training options, test options, and basic options (used in both training and test)."""

View File

@@ -0,0 +1,138 @@
import argparse
import os
from util import util
import torch
import models
import data
class BaseOptions():
"""This class defines options used during both training and test time.
It also implements several helper functions such as parsing, printing, and saving the options.
It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.
"""
def __init__(self):
"""Reset the class; indicates the class hasn't been initailized"""
self.initialized = False
def initialize(self, parser):
"""Define the common options that are used in both training and test."""
# basic parameters
parser.add_argument('--dataroot', type=str, default='./datasets/ITS_v2', help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
# model parameters
parser.add_argument('--model', type=str, default='cycle_gan', help='chooses which model to use. [cycle_gan | pix2pix | test | colorization]')
parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')
parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels: 3 for RGB and 1 for grayscale')
parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in the last conv layer')
parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in the first conv layer')
parser.add_argument('--netD', type=str, default='basic', help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')
parser.add_argument('--netG', type=str, default='resnet_9blocks', help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')
parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')
parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization [instance | batch | none]')
parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')
parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')
# dataset parameters
parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single | colorization]')
parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA')
parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
parser.add_argument('--load_size', type=int, default=286, help='scale images to this size')
parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')
parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')
parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
parser.add_argument('--display_winsize', type=int, default=256, help='display window size for both visdom and HTML')
# additional parameters
parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')
parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')
self.initialized = True
return parser
def gather_options(self, asigned_parser=None):
"""Initialize our parser with basic options(only once).
Add additional model-specific and dataset-specific options.
These options are defined in the <modify_commandline_options> function
in model and dataset classes.
"""
if not self.initialized: # check if it has been initialized
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser = self.initialize(parser)
else:
parser = asigned_parser
# get the basic options
opt, _ = parser.parse_known_args()
# modify model-related parser options
model_name = opt.model
model_option_setter = models.get_option_setter(model_name)
parser = model_option_setter(parser, self.isTrain)
opt, _ = parser.parse_known_args() # parse again with new defaults
# modify dataset-related parser options
dataset_name = opt.dataset_mode
dataset_option_setter = data.get_option_setter(dataset_name)
parser = dataset_option_setter(parser, self.isTrain)
# save and return the parser
self.parser = parser
return parser.parse_args()
def print_options(self, opt):
"""Print and save options
It will print both current options and default values(if different).
It will save options into a text file / [checkpoints_dir] / opt.txt
"""
message = ''
message += '----------------- Options ---------------\n'
for k, v in sorted(vars(opt).items()):
comment = ''
default = self.parser.get_default(k)
if v != default:
comment = '\t[default: %s]' % str(default)
message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
message += '----------------- End -------------------'
print(message)
# save to the disk
expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
util.mkdirs(expr_dir)
file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
with open(file_name, 'wt') as opt_file:
opt_file.write(message)
opt_file.write('\n')
def parse(self, asigned_parser=None):
"""Parse our options, create checkpoints directory suffix, and set up gpu device."""
opt = self.gather_options(asigned_parser)
opt.isTrain = self.isTrain # train or test
# process opt.suffix
if opt.suffix:
suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
opt.name = opt.name + suffix
self.print_options(opt)
# set gpu ids
str_ids = opt.gpu_ids.split(',')
opt.gpu_ids = []
for str_id in str_ids:
id = int(str_id)
if id >= 0:
opt.gpu_ids.append(id)
if len(opt.gpu_ids) > 0:
torch.cuda.set_device(opt.gpu_ids[0])
self.opt = opt
return self.opt

View File

@@ -0,0 +1,27 @@
from .base_options import BaseOptions
class TestOptions(BaseOptions):
"""This class includes test options.
It also includes shared options defined in BaseOptions.
"""
def initialize(self, parser):
parser = BaseOptions.initialize(self, parser) # define shared options
parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
# Dropout and Batchnorm has different behavioir during training and test.
parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
parser.add_argument('--num_test', type=int, default=50, help='how many test images to run')
parser.add_argument('--save_image', action='store_true', help='save result images.')
parser.add_argument('--method_name', type=str, default='Mine', help='short name for your dehazing method')
# rewrite devalue values
parser.set_defaults(model='test')
# To avoid cropping, the load_size should be the same as crop_size
parser.set_defaults(load_size=parser.get_default('crop_size'))
self.isTrain = False
return parser

View File

@@ -0,0 +1,40 @@
from .base_options import BaseOptions
class TrainOptions(BaseOptions):
"""This class includes training options.
It also includes shared options defined in BaseOptions.
"""
def initialize(self, parser):
parser = BaseOptions.initialize(self, parser)
# visdom and HTML visualization parameters
parser.add_argument('--display_freq', type=int, default=400, help='frequency of showing training results on screen')
parser.add_argument('--display_ncols', type=int, default=4, help='if positive, display all images in a single visdom web panel with certain number of images per row.')
parser.add_argument('--display_id', type=int, default=1, help='window id of the web display')
parser.add_argument('--display_server', type=str, default="http://localhost", help='visdom server of the web display')
parser.add_argument('--display_env', type=str, default='main', help='visdom display environment name (default is "main")')
parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')
parser.add_argument('--update_html_freq', type=int, default=1000, help='frequency of saving training results to html')
parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
# network saving and loading parameters
parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')
parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')
parser.add_argument('--save_by_iter', action='store_true', help='whether saves model by iteration')
parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
# training parameters
parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate')
parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
parser.add_argument('--gan_mode', type=str, default='lsgan', help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')
parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')
parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy. [linear | step | plateau | cosine]')
parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')
self.isTrain = True
return parser

77
RefineDNet/quick_test.py Normal file
View File

@@ -0,0 +1,77 @@
import os,time
import ntpath
import numpy as np
import scipy.io as sio
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
from util import util
if __name__ == '__main__':
opt = TestOptions().parse() # get test options
# hard-code some parameters for test
opt.num_threads = 0 # test code only supports num_threads = 1
opt.batch_size = 1 # test code only supports batch_size = 1
opt.serial_batches = True # disable data shuffling; comment this line if results on randomly chosen images are needed.
opt.no_flip = True # no flip; comment this line if results on flipped images are needed.
opt.display_id = -1 # no visdom display; the test code saves the results to a HTML file.
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
model = create_model(opt) # create a model given opt.model and other options
model.setup(opt) # regular setup: load and print networks; create schedulers
if opt.save_image:
curSaveFolder = os.path.join(opt.dataroot, opt.method_name)
if not os.path.exists(curSaveFolder):
os.makedirs(curSaveFolder, mode=0o777)
# test with eval mode. This only affects layers like batchnorm and dropout.
# For [pix2pix]: we use batchnorm and dropout in the original pix2pix. You can experiment it with and without eval() mode.
# For [CycleGAN]: It should not affect CycleGAN as CycleGAN uses instancenorm without dropout.
if opt.eval:
model.eval()
time_total = 0
for i, data in enumerate(dataset):
# if i <= 627:
# continue
img_path = data['paths']
short_path = ntpath.basename(img_path[0])
name = os.path.splitext(short_path)[0]
print('%s [%d]'%(short_path, i+1))
# print(data['B_paths'])
if 'haze' in data.keys():
minSize = min(data['haze'].shape[2:4])
else:
minSize = min(data['A'].shape[2:4])
if minSize < 256:
print(' skip because the minimum size is %s'%minSize)
continue
# if i >= opt.num_test: # only apply our model to opt.num_test images.
# break
t0 = time.time()
model.set_input(data) # unpack data from data loader
model.test() # run inference
time_total += time.time() - t0
visuals = model.get_current_visuals() # get image results
rec_J = util.tensor2im(visuals['rec_J'], float)/255. # [0, 1]
refine_J = util.tensor2im(visuals['refine_J'], float)/255. # [0, 1]
real_I = util.tensor2im(data['haze'], float) # [0, 255], float
result_J = util.fuse_images(real_I, rec_J*255., refine_J*255.)/255. # [0, 1], np.float
# save result images
if opt.save_image:
dehzImg = (result_J*255).astype(np.uint8) #[0, 255], np.uint8
util.save_image(dehzImg, os.path.join(curSaveFolder, '%s_dehz.png'%(name)))
# refinedT = util.tensor2im(visuals['refine_T_vis'])
# util.save_image(refinedT, os.path.join(curSaveFolder, '%s_ref_T.png'%(name)))
print('num: %d'%len(dataset))
print('average time: %f'%(time_total/len(dataset)))

51
RefineDNet/test_BeDDE.py Normal file
View File

@@ -0,0 +1,51 @@
import os, ntpath
import numpy as np
import scipy.io as sio
import torchvision.utils as vutils
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
from util import util
if __name__ == '__main__':
opt = TestOptions().parse() # get test options
opt.nThreads = 1 # mytest code only supports nThreads = 1
opt.batchSize = 1 # mytest code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
model = create_model(opt) # create a model given opt.model and other options
model.setup(opt) # regular setup: load and print networks; create schedulers
if opt.eval:
model.eval()
for i, data in enumerate(dataset):
model.set_input(data) # unpack data from data loader
model.test() # run inference
visuals = model.get_current_visuals() # get image results
real_I = util.tensor2im(data['haze'], float) # [0, 255], float
real_J = util.tensor2im(data['clear'], float) # [0, 255], float
rec_J = util.tensor2im(visuals['rec_J'], float) # [0, 255], float
refine_J = util.tensor2im(visuals['refine_J'], float) # [0, 255], float
result_J = util.fuse_images(real_I, rec_J, refine_J) # [0, 255], np.float
img_paths = model.get_image_paths() # get image paths
short_path = ntpath.basename(img_paths[0])
name = os.path.splitext(short_path)[0]
print('processing image %s (%d/%d)'%(short_path, i+1, len(dataset)))
if opt.save_image:
curSaveFolder = os.path.join(opt.dataroot, data['city'][0], opt.method_name)
if not os.path.exists(curSaveFolder):
os.makedirs(curSaveFolder, mode=0o777)
dehzImg = (result_J).astype(np.uint8) #[0, 255], np.uint8
util.save_image(dehzImg, os.path.join(curSaveFolder, '%s_%s.png'%(name, opt.method_name)))

77
RefineDNet/train.py Normal file
View File

@@ -0,0 +1,77 @@
"""General-purpose training script for image-to-image translation.
This script works for various models (with option '--model': e.g., pix2pix, cyclegan, colorization) and
different datasets (with option '--dataset_mode': e.g., aligned, unaligned, single, colorization).
You need to specify the dataset ('--dataroot'), experiment name ('--name'), and model ('--model').
It first creates model, dataset, and visualizer given the option.
It then does standard network training. During the training, it also visualize/save the images, print/save the loss plot, and save models.
The script supports continue/resume training. Use '--continue_train' to resume your previous training.
Example:
Train a CycleGAN model:
python train.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan
Train a pix2pix model:
python train.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA
See options/base_options.py and options/train_options.py for more training options.
See training and test tips at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md
See frequently asked questions at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md
"""
import time
from options.train_options import TrainOptions
from data import create_dataset
from models import create_model
from util.visualizer import Visualizer
if __name__ == '__main__':
opt = TrainOptions().parse() # get training options
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
dataset_size = len(dataset) # get the number of images in the dataset.
print('The number of training images = %d' % dataset_size)
model = create_model(opt) # create a model given opt.model and other options
model.setup(opt) # regular setup: load and print networks; create schedulers
visualizer = Visualizer(opt) # create a visualizer that display/save images and plots
total_iters = 0 # the total number of training iterations
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1): # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>
epoch_start_time = time.time() # timer for entire epoch
iter_data_time = time.time() # timer for data loading per iteration
epoch_iter = 0 # the number of training iterations in current epoch, reset to 0 every epoch
for i, data in enumerate(dataset): # inner loop within one epoch
iter_start_time = time.time() # timer for computation per iteration
if total_iters % opt.print_freq == 0:
t_data = iter_start_time - iter_data_time
visualizer.reset()
total_iters += opt.batch_size
epoch_iter += opt.batch_size
model.set_input(data) # unpack data from dataset and apply preprocessing
model.optimize_parameters() # calculate loss functions, get gradients, update network weights
if total_iters % opt.display_freq == 0: # display images on visdom and save images to a HTML file
save_result = total_iters % opt.update_html_freq == 0
model.compute_visuals()
visualizer.display_current_results(model.get_current_visuals(), epoch, save_result)
if total_iters % opt.print_freq == 0: # print training losses and save logging information to the disk
losses = model.get_current_losses()
t_comp = (time.time() - iter_start_time) / opt.batch_size
visualizer.print_current_losses(epoch, epoch_iter, losses, t_comp, t_data)
if opt.display_id > 0:
visualizer.plot_current_losses(epoch, float(epoch_iter) / dataset_size, losses)
if total_iters % opt.save_latest_freq == 0: # cache our latest model every <save_latest_freq> iterations
print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters))
save_suffix = 'iter_%d' % total_iters if opt.save_by_iter else 'latest'
model.save_networks(save_suffix)
iter_data_time = time.time()
if epoch % opt.save_epoch_freq == 0: # cache our model every <save_epoch_freq> epochs
print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters))
model.save_networks('latest')
model.save_networks(epoch)
print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
model.update_learning_rate() # update learning rates at the end of every epoch.

View File

@@ -0,0 +1 @@
"""This package includes a miscellaneous collection of useful helper functions."""

110
RefineDNet/util/get_data.py Normal file
View File

@@ -0,0 +1,110 @@
from __future__ import print_function
import os
import tarfile
import requests
from warnings import warn
from zipfile import ZipFile
from bs4 import BeautifulSoup
from os.path import abspath, isdir, join, basename
class GetData(object):
"""A Python script for downloading CycleGAN or pix2pix datasets.
Parameters:
technique (str) -- One of: 'cyclegan' or 'pix2pix'.
verbose (bool) -- If True, print additional information.
Examples:
>>> from util.get_data import GetData
>>> gd = GetData(technique='cyclegan')
>>> new_data_path = gd.get(save_path='./datasets') # options will be displayed.
Alternatively, You can use bash scripts: 'scripts/download_pix2pix_model.sh'
and 'scripts/download_cyclegan_model.sh'.
"""
def __init__(self, technique='cyclegan', verbose=True):
url_dict = {
'pix2pix': 'http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/',
'cyclegan': 'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets'
}
self.url = url_dict.get(technique.lower())
self._verbose = verbose
def _print(self, text):
if self._verbose:
print(text)
@staticmethod
def _get_options(r):
soup = BeautifulSoup(r.text, 'lxml')
options = [h.text for h in soup.find_all('a', href=True)
if h.text.endswith(('.zip', 'tar.gz'))]
return options
def _present_options(self):
r = requests.get(self.url)
options = self._get_options(r)
print('Options:\n')
for i, o in enumerate(options):
print("{0}: {1}".format(i, o))
choice = input("\nPlease enter the number of the "
"dataset above you wish to download:")
return options[int(choice)]
def _download_data(self, dataset_url, save_path):
if not isdir(save_path):
os.makedirs(save_path)
base = basename(dataset_url)
temp_save_path = join(save_path, base)
with open(temp_save_path, "wb") as f:
r = requests.get(dataset_url)
f.write(r.content)
if base.endswith('.tar.gz'):
obj = tarfile.open(temp_save_path)
elif base.endswith('.zip'):
obj = ZipFile(temp_save_path, 'r')
else:
raise ValueError("Unknown File Type: {0}.".format(base))
self._print("Unpacking Data...")
obj.extractall(save_path)
obj.close()
os.remove(temp_save_path)
def get(self, save_path, dataset=None):
"""
Download a dataset.
Parameters:
save_path (str) -- A directory to save the data to.
dataset (str) -- (optional). A specific dataset to download.
Note: this must include the file extension.
If None, options will be presented for you
to choose from.
Returns:
save_path_full (str) -- the absolute path to the downloaded data.
"""
if dataset is None:
selected_dataset = self._present_options()
else:
selected_dataset = dataset
save_path_full = join(save_path, selected_dataset.split('.')[0])
if isdir(save_path_full):
warn("\n'{0}' already exists. Voiding Download.".format(
save_path_full))
else:
self._print('Downloading Data...')
url = "{0}/{1}".format(self.url, selected_dataset)
self._download_data(url, save_path=save_path)
return abspath(save_path_full)

86
RefineDNet/util/html.py Normal file
View File

@@ -0,0 +1,86 @@
import dominate
from dominate.tags import meta, h3, table, tr, td, p, a, img, br
import os
class HTML:
"""This HTML class allows us to save images and write texts into a single HTML file.
It consists of functions such as <add_header> (add a text header to the HTML file),
<add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
It is based on Python library 'dominate', a Python library for creating and manipulating HTML documents using a DOM API.
"""
def __init__(self, web_dir, title, refresh=0):
"""Initialize the HTML classes
Parameters:
web_dir (str) -- a directory that stores the webpage. HTML file will be created at <web_dir>/index.html; images will be saved at <web_dir/images/
title (str) -- the webpage name
refresh (int) -- how often the website refresh itself; if 0; no refreshing
"""
self.title = title
self.web_dir = web_dir
self.img_dir = os.path.join(self.web_dir, 'images')
if not os.path.exists(self.web_dir):
os.makedirs(self.web_dir)
if not os.path.exists(self.img_dir):
os.makedirs(self.img_dir)
self.doc = dominate.document(title=title)
if refresh > 0:
with self.doc.head:
meta(http_equiv="refresh", content=str(refresh))
def get_image_dir(self):
"""Return the directory that stores images"""
return self.img_dir
def add_header(self, text):
"""Insert a header to the HTML file
Parameters:
text (str) -- the header text
"""
with self.doc:
h3(text)
def add_images(self, ims, txts, links, width=400):
"""add images to the HTML file
Parameters:
ims (str list) -- a list of image paths
txts (str list) -- a list of image names shown on the website
links (str list) -- a list of hyperref links; when you click an image, it will redirect you to a new page
"""
self.t = table(border=1, style="table-layout: fixed;") # Insert a table
self.doc.add(self.t)
with self.t:
with tr():
for im, txt, link in zip(ims, txts, links):
with td(style="word-wrap: break-word;", halign="center", valign="top"):
with p():
with a(href=os.path.join('images', link)):
img(style="width:%dpx" % width, src=os.path.join('images', im))
br()
p(txt)
def save(self):
"""save the current content to the HMTL file"""
html_file = '%s/index.html' % self.web_dir
f = open(html_file, 'wt')
f.write(self.doc.render())
f.close()
if __name__ == '__main__': # we show an example usage here.
html = HTML('web/', 'test_html')
html.add_header('hello world')
ims, txts, links = [], [], []
for n in range(4):
ims.append('image_%d.png' % n)
txts.append('text_%d' % n)
links.append('image_%d.png' % n)
html.add_images(ims, txts, links)
html.save()

View File

@@ -0,0 +1,54 @@
import random
import torch
class ImagePool():
"""This class implements an image buffer that stores previously generated images.
This buffer enables us to update discriminators using a history of generated images
rather than the ones produced by the latest generators.
"""
def __init__(self, pool_size):
"""Initialize the ImagePool class
Parameters:
pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
"""
self.pool_size = pool_size
if self.pool_size > 0: # create an empty pool
self.num_imgs = 0
self.images = []
def query(self, images):
"""Return an image from the pool.
Parameters:
images: the latest generated images from the generator
Returns images from the buffer.
By 50/100, the buffer will return input images.
By 50/100, the buffer will return images previously stored in the buffer,
and insert the current images to the buffer.
"""
if self.pool_size == 0: # if the buffer size is 0, do nothing
return images
return_images = []
for image in images:
image = torch.unsqueeze(image.data, 0)
if self.num_imgs < self.pool_size: # if the buffer is not full; keep inserting current images to the buffer
self.num_imgs = self.num_imgs + 1
self.images.append(image)
return_images.append(image)
else:
p = random.uniform(0, 1)
if p > 0.5: # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
random_id = random.randint(0, self.pool_size - 1) # randint is inclusive
tmp = self.images[random_id].clone()
self.images[random_id] = image
return_images.append(tmp)
else: # by another 50% chance, the buffer will return the current image
return_images.append(image)
return_images = torch.cat(return_images, 0) # collect all the images and return
return return_images

271
RefineDNet/util/util.py Normal file
View File

@@ -0,0 +1,271 @@
"""This module contains simple helper functions """
from __future__ import print_function
import torch
import numpy as np
from PIL import Image
import scipy.ndimage as ndimage
import os
import cv2
import torch.nn.functional as F
def synthesize_fog(J, t, A=None):
"""
Synthesize hazy image base on optical model
I = J * t + A * (1 - t)
"""
if A is None:
A = 1
return J * t + A * (1 - t)
def reverse_fog(I, t, A=1, t0=0.01):
"""
Recover haze-free image using hazy image and depth
J = (I - A) / max(t, t0) + A
"""
t_clamp = torch.clamp(t, t0, 1)
J = (I-A) / t_clamp + A
return torch.clamp(J, -1, 1)
def fuse_images(real_I, rec_J, refine_J):
"""
real_I, rec_J, and refine_J: Images with shape hxwx3
"""
# realness features
mat_RGB2YMN = np.array([[0.299,0.587,0.114],
[0.30,0.04,-0.35],
[0.34,-0.6,0.17]])
recH,recW,recChl = rec_J.shape
rec_J_flat = rec_J.reshape([recH*recW,recChl])
rec_J_flat_YMN = (mat_RGB2YMN.dot(rec_J_flat.T)).T
rec_J_YMN = rec_J_flat_YMN.reshape(rec_J.shape)
refine_J_flat = refine_J.reshape([recH*recW,recChl])
refine_J_flat_YMN = (mat_RGB2YMN.dot(refine_J_flat.T)).T
refine_J_YMN = refine_J_flat_YMN.reshape(refine_J.shape)
real_I_flat = real_I.reshape([recH*recW,recChl])
real_I_flat_YMN = (mat_RGB2YMN.dot(real_I_flat.T)).T
real_I_YMN = real_I_flat_YMN.reshape(real_I.shape)
# gradient features
rec_Gx = cv2.Sobel(rec_J_YMN[:,:,0],cv2.CV_64F,1,0,ksize=3)
rec_Gy = cv2.Sobel(rec_J_YMN[:,:,0],cv2.CV_64F,0,1,ksize=3)
rec_GM = np.sqrt(rec_Gx**2 + rec_Gy**2)
refine_Gx = cv2.Sobel(refine_J_YMN[:,:,0],cv2.CV_64F,1,0,ksize=3)
refine_Gy = cv2.Sobel(refine_J_YMN[:,:,0],cv2.CV_64F,0,1,ksize=3)
refine_GM = np.sqrt(refine_Gx**2 + refine_Gy**2)
real_Gx = cv2.Sobel(real_I_YMN[:,:,0],cv2.CV_64F,1,0,ksize=3)
real_Gy = cv2.Sobel(real_I_YMN[:,:,0],cv2.CV_64F,0,1,ksize=3)
real_GM = np.sqrt(real_Gx**2 + real_Gy**2)
# similarity
rec_S_V = (2*real_GM*rec_GM+160)/(real_GM**2+rec_GM**2+160)
rec_S_M = (2*rec_J_YMN[:,:,1]*real_I_YMN[:,:,1]+130)/(rec_J_YMN[:,:,1]**2+real_I_YMN[:,:,1]**2+130)
rec_S_N = (2*rec_J_YMN[:,:,2]*real_I_YMN[:,:,2]+130)/(rec_J_YMN[:,:,2]**2+real_I_YMN[:,:,2]**2+130)
rec_S_R = (rec_S_M*rec_S_N).reshape([recH,recW])
refine_S_V = (2*real_GM*refine_GM+160)/(real_GM**2+refine_GM**2+160)
refine_S_M = (2*refine_J_YMN[:,:,1]*real_I_YMN[:,:,1]+130)/(refine_J_YMN[:,:,1]**2+real_I_YMN[:,:,1]**2+130)
refine_S_N = (2*refine_J_YMN[:,:,2]*real_I_YMN[:,:,2]+130)/(refine_J_YMN[:,:,2]**2+real_I_YMN[:,:,2]**2+130)
refine_S_R = (refine_S_M*refine_S_N).reshape([recH,recW])
rec_S = rec_S_R*np.power(rec_S_V, 0.4)
refine_S = refine_S_R*np.power(refine_S_V, 0.4)
fuseWeight = np.exp(rec_S)/(np.exp(rec_S)+np.exp(refine_S))
fuseWeightMap = fuseWeight.reshape([recH,recW,1]).repeat(3,axis=2)
fuse_J = rec_J*fuseWeightMap + refine_J*(1-fuseWeightMap)
return fuse_J
def get_tensor_dark_channel(img, neighborhood_size):
shape = img.shape
if len(shape) == 4:
img_min = torch.min(img, dim=1)
img_dark = F.max_pool2d(img_min, kernel_size=neighborhood_size, stride=1)
else:
raise NotImplementedError('get_tensor_dark_channel is only for 4-d tensor [N*C*H*W]')
return img_dark
def array2Tensor(in_array, gpu_id=-1):
in_shape = in_array.shape
if len(in_shape) == 2:
in_array = in_array[:,:,np.newaxis]
arr_tmp = in_array.transpose([2,0,1])
arr_tmp = arr_tmp[np.newaxis,:]
if gpu_id >= 0:
return torch.tensor(arr_tmp.astype(float)).to(gpu_id)
else:
return torch.tensor(arr_tmp.astype(float))
def tensor2im(input_image, imtype=np.uint8):
""""Converts a Tensor array into a numpy image array.
Parameters:
input_image (tensor) -- the input image tensor array
imtype (type) -- the desired type of the converted numpy array
"""
if not isinstance(input_image, np.ndarray):
if isinstance(input_image, torch.Tensor): # get the data from a variable
image_tensor = input_image.data
else:
return input_image
image_numpy = image_tensor[0].cpu().float().numpy() # convert it into a numpy array
if image_numpy.shape[0] == 1: # grayscale to RGB
image_numpy = np.tile(image_numpy, (3, 1, 1))
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 # post-processing: tranpose and scaling
else: # if it is a numpy array, do nothing
image_numpy = input_image
return image_numpy.astype(imtype)
def rescale_tensor(input_tensor):
""""Converts a Tensor array into the Tensor array whose data are identical to the image's.
[height, width] not [width, height]
Parameters:
input_image (tensor) -- the input image tensor array
imtype (type) -- the desired type of the converted numpy array
"""
if isinstance(input_tensor, torch.Tensor):
input_tmp = input_tensor.cpu().float()
output_tmp = (input_tmp + 1) / 2.0 * 255.0
output_tmp = output_tmp.to(torch.uint8)
else:
return input_tensor
return output_tmp.to(torch.float32) / 255.0
# if not isinstance(input_image, np.ndarray):
# if isinstance(input_image, torch.Tensor): # get the data from a variable
# image_tensor = input_image.data
# else:
# return input_image
# image_numpy = image_tensor.cpu().float().numpy() # convert it into a numpy array
# image_numpy = (image_numpy + 1) / 2.0 * white_color # post-processing: tranpose and scaling
# else: # if it is a numpy array, do nothing
# image_numpy = input_image
# return torch.from_numpy(image_numpy)
def my_imresize(in_array, tar_size):
oh = in_array.shape[0]
ow = in_array.shape[1]
if len(tar_size) == 2:
h_ratio = tar_size[0]/oh
w_ratio = tar_size[1]/ow
elif len(tar_size) == 1:
h_ratio = tar_size
w_ratio = tar_size
if len(in_array.shape) == 3:
return ndimage.zoom(in_array, (h_ratio, w_ratio, 1), prefilter=False)
else:
return ndimage.zoom(in_array, (h_ratio, w_ratio), prefilter=False)
def psnr(img, ref, max_val=1):
if isinstance(img, torch.Tensor):
distImg = img.cpu().float().numpy()
elif isinstance(img, np.ndarray):
distImg = img.astype(float)
else:
distImg = np.array(img).astype(float)
if isinstance(ref, torch.Tensor):
refImg = ref.cpu().float().numpy()
elif isinstance(ref, np.ndarray):
refImg = ref.astype(float)
else:
refImg = np.array(ref).astype(float)
rmse = np.sqrt( ((distImg-refImg)**2).mean() )
# rmse = np.std(distImg-refImg) # keep the same with RESIDE's criterion
return 20*np.log10(max_val/rmse)
def diagnose_network(net, name='network'):
"""Calculate and print the mean of average absolute(gradients)
Parameters:
net (torch network) -- Torch network
name (str) -- the name of the network
"""
mean = 0.0
count = 0
for param in net.parameters():
if param.grad is not None:
mean += torch.mean(torch.abs(param.grad.data))
count += 1
if count > 0:
mean = mean / count
print(name)
print(mean)
def save_image(image_numpy, image_path):
"""Save a numpy image to the disk
Parameters:
image_numpy (numpy array) -- input numpy array
image_path (str) -- the path of the image
"""
image_pil = Image.fromarray(image_numpy)
image_pil.save(image_path)
def print_numpy(x, val=True, shp=False):
"""Print the mean, min, max, median, std, and size of a numpy array
Parameters:
val (bool) -- if print the values of the numpy array
shp (bool) -- if print the shape of the numpy array
"""
x = x.astype(np.float64)
if shp:
print('shape,', x.shape)
if val:
x = x.flatten()
print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
def mkdirs(paths):
"""create empty directories if they don't exist
Parameters:
paths (str list) -- a list of directory paths
"""
if isinstance(paths, list) and not isinstance(paths, str):
for path in paths:
mkdir(path)
else:
mkdir(paths)
def mkdir(path):
"""create a single empty directory if it didn't exist
Parameters:
path (str) -- a single directory path
"""
if not os.path.exists(path):
os.makedirs(path)

View File

@@ -0,0 +1,227 @@
import numpy as np
import os
import sys
import ntpath
import time
from . import util, html
from subprocess import Popen, PIPE
from scipy.misc import imresize
if sys.version_info[0] == 2:
VisdomExceptionBase = Exception
else:
VisdomExceptionBase = ConnectionError
def save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=256):
"""Save images to the disk.
Parameters:
webpage (the HTML class) -- the HTML webpage class that stores these imaegs (see html.py for more details)
visuals (OrderedDict) -- an ordered dictionary that stores (name, images (either tensor or numpy) ) pairs
image_path (str) -- the string is used to create image paths
aspect_ratio (float) -- the aspect ratio of saved images
width (int) -- the images will be resized to width x width
This function will save images stored in 'visuals' to the HTML file specified by 'webpage'.
"""
image_dir = webpage.get_image_dir()
short_path = ntpath.basename(image_path[0])
name = os.path.splitext(short_path)[0]
webpage.add_header(name)
ims, txts, links = [], [], []
for label, im_data in visuals.items():
im = util.tensor2im(im_data)
image_name = '%s_%s.png' % (name, label)
save_path = os.path.join(image_dir, image_name)
h, w, _ = im.shape
if aspect_ratio > 1.0:
im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic')
if aspect_ratio < 1.0:
im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic')
util.save_image(im, save_path)
ims.append(image_name)
txts.append(label)
links.append(image_name)
webpage.add_images(ims, txts, links, width=width)
class Visualizer():
"""This class includes several functions that can display/save images and print/save logging information.
It uses a Python library 'visdom' for display, and a Python library 'dominate' (wrapped in 'HTML') for creating HTML files with images.
"""
def __init__(self, opt):
"""Initialize the Visualizer class
Parameters:
opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
Step 1: Cache the training/test options
Step 2: connect to a visdom server
Step 3: create an HTML object for saveing HTML filters
Step 4: create a logging file to store training losses
"""
self.opt = opt # cache the option
self.display_id = opt.display_id
self.use_html = opt.isTrain and not opt.no_html
self.win_size = opt.display_winsize
self.name = opt.name
self.port = opt.display_port
self.saved = False
if self.display_id > 0: # connect to a visdom server given <display_port> and <display_server>
import visdom
self.ncols = opt.display_ncols
self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port, env=opt.display_env)
if not self.vis.check_connection():
self.create_visdom_connections()
if self.use_html: # create an HTML object at <checkpoints_dir>/web/; images will be saved under <checkpoints_dir>/web/images/
self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
self.img_dir = os.path.join(self.web_dir, 'images')
print('create web directory %s...' % self.web_dir)
util.mkdirs([self.web_dir, self.img_dir])
# create a logging file to store training losses
self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
with open(self.log_name, "a") as log_file:
now = time.strftime("%c")
log_file.write('================ Training Loss (%s) ================\n' % now)
def reset(self):
"""Reset the self.saved status"""
self.saved = False
def create_visdom_connections(self):
"""If the program could not connect to Visdom server, this function will start a new server at port < self.port > """
cmd = sys.executable + ' -m visdom.server -p %d &>/dev/null &' % self.port
print('\n\nCould not connect to Visdom server. \n Trying to start a server....')
print('Command: %s' % cmd)
Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE)
def display_current_results(self, visuals, epoch, save_result):
"""Display current results on visdom; save current results to an HTML file.
Parameters:
visuals (OrderedDict) - - dictionary of images to display or save
epoch (int) - - the current epoch
save_result (bool) - - if save the current results to an HTML file
"""
if self.display_id > 0: # show images in the browser using visdom
ncols = self.ncols
if ncols > 0: # show all the images in one visdom panel
ncols = min(ncols, len(visuals))
h, w = next(iter(visuals.values())).shape[:2]
table_css = """<style>
table {border-collapse: separate; border-spacing: 4px; white-space: nowrap; text-align: center}
table td {width: % dpx; height: % dpx; padding: 4px; outline: 4px solid black}
</style>""" % (w, h) # create a table css
# create a table of images.
title = self.name
label_html = ''
label_html_row = ''
images = []
idx = 0
for label, image in visuals.items():
image_numpy = util.tensor2im(image)
label_html_row += '<td>%s</td>' % label
images.append(image_numpy.transpose([2, 0, 1]))
idx += 1
if idx % ncols == 0:
label_html += '<tr>%s</tr>' % label_html_row
label_html_row = ''
white_image = np.ones_like(image_numpy.transpose([2, 0, 1])) * 255
while idx % ncols != 0:
images.append(white_image)
label_html_row += '<td></td>'
idx += 1
if label_html_row != '':
label_html += '<tr>%s</tr>' % label_html_row
try:
self.vis.images(images, nrow=ncols, win=self.display_id + 1,
padding=2, opts=dict(title=title + ' images'))
label_html = '<table>%s</table>' % label_html
self.vis.text(table_css + label_html, win=self.display_id + 2,
opts=dict(title=title + ' labels'))
except VisdomExceptionBase:
self.create_visdom_connections()
else: # show each image in a separate visdom panel;
idx = 1
try:
for label, image in visuals.items():
image_numpy = util.tensor2im(image)
self.vis.image(image_numpy.transpose([2, 0, 1]), opts=dict(title=label),
win=self.display_id + idx)
idx += 1
except VisdomExceptionBase:
self.create_visdom_connections()
if self.use_html and (save_result or not self.saved): # save images to an HTML file if they haven't been saved.
self.saved = True
# save images to the disk
for label, image in visuals.items():
image_numpy = util.tensor2im(image)
img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label))
util.save_image(image_numpy, img_path)
# update website
webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=1)
for n in range(epoch, 0, -1):
webpage.add_header('epoch [%d]' % n)
ims, txts, links = [], [], []
for label, image_numpy in visuals.items():
image_numpy = util.tensor2im(image)
img_path = 'epoch%.3d_%s.png' % (n, label)
ims.append(img_path)
txts.append(label)
links.append(img_path)
webpage.add_images(ims, txts, links, width=self.win_size)
webpage.save()
def plot_current_losses(self, epoch, counter_ratio, losses):
"""display the current losses on visdom display: dictionary of error labels and values
Parameters:
epoch (int) -- current epoch
counter_ratio (float) -- progress (percentage) in the current epoch, between 0 to 1
losses (OrderedDict) -- training losses stored in the format of (name, float) pairs
"""
if not hasattr(self, 'plot_data'):
self.plot_data = {'X': [], 'Y': [], 'legend': list(losses.keys())}
self.plot_data['X'].append(epoch + counter_ratio)
self.plot_data['Y'].append([losses[k] for k in self.plot_data['legend']])
try:
self.vis.line(
X=np.stack([np.array(self.plot_data['X'])] * len(self.plot_data['legend']), 1),
Y=np.array(self.plot_data['Y']),
opts={
'title': self.name + ' loss over time',
'legend': self.plot_data['legend'],
'xlabel': 'epoch',
'ylabel': 'loss'},
win=self.display_id)
except VisdomExceptionBase:
self.create_visdom_connections()
# losses: same format as |losses| of plot_current_losses
def print_current_losses(self, epoch, iters, losses, t_comp, t_data):
"""print current losses on console; also save the losses to the disk
Parameters:
epoch (int) -- current epoch
iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch)
losses (OrderedDict) -- training losses stored in the format of (name, float) pairs
t_comp (float) -- computational time per data point (normalized by batch_size)
t_data (float) -- data loading time per data point (normalized by batch_size)
"""
message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, iters, t_comp, t_data)
for k, v in losses.items():
message += '%s: %.3f ' % (k, v)
print(message) # print the message
with open(self.log_name, "a") as log_file:
log_file.write('%s\n' % message) # save the message

View File

@@ -0,0 +1,4 @@
# Copy to config/baidu_api.env and fill values locally.
# The real config/baidu_api.env file is ignored by git.
BAIDU_API_KEY=your_api_key
BAIDU_SECRET_KEY=your_secret_key

8
run_dehaze_web.sh Executable file
View File

@@ -0,0 +1,8 @@
#!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")"
export DEHAZE_CAFFE_PYTHON="${DEHAZE_CAFFE_PYTHON:-/home/wkmgc/miniconda3/envs/dehaze_caffe/bin/python}"
export DEHAZE_TORCH_PYTHON="${DEHAZE_TORCH_PYTHON:-/home/wkmgc/miniconda3/envs/seg_server/bin/python}"
python web_dehaze/server.py --host 0.0.0.0 --port 7860

12
scripts/clean_generated.sh Executable file
View File

@@ -0,0 +1,12 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"
find . -type d -name '__pycache__' -prune -exec rm -rf {} +
find . -type f \( -name '*.pyc' -o -name '*.pyo' -o -name 'Thumbs.db' -o -name '.DS_Store' \) -delete
rm -rf web_results
mkdir -p web_results
echo "Cleaned generated files."

58
scripts/verify_all.py Executable file
View File

@@ -0,0 +1,58 @@
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "web_dehaze"))
import pipeline # noqa: E402
DEFAULT_METHODS = ["AOD", "Baidu_API", "DCP", "DehazeNet", "GCANet", "RefineDNet"]
def main() -> int:
parser = argparse.ArgumentParser(description="Verify dehaze models through the unified pipeline.")
parser.add_argument("--images", nargs="*", help="Image names under 待去雾图片/. Defaults to all images.")
parser.add_argument("--methods", nargs="*", default=DEFAULT_METHODS, help="Methods to run.")
parser.add_argument("--skip-baidu", action="store_true", help="Skip Baidu API calls.")
args = parser.parse_args()
images = args.images or [item["name"] for item in pipeline.list_images()]
methods = [m for m in args.methods if not (args.skip_baidu and m == "Baidu_API")]
caps = pipeline.capabilities()["methods"]
unavailable = [m for m in methods if not caps.get(m, {}).get("available")]
if unavailable:
print("Unavailable methods:", ", ".join(unavailable))
return 2
failures: list[tuple[str, str, str]] = []
for image in images:
print(f"\n=== {image} ===")
for method in methods:
start = time.time()
print(f"[RUN] {method}")
try:
result = pipeline.run_dehaze_method(image, method, {}, print)
elapsed = time.time() - start
print(f"[OK] {method}: {pipeline.relpath(result)} ({elapsed:.1f}s)")
except Exception as exc:
failures.append((image, method, str(exc)))
print(f"[FAIL] {method}: {exc}")
if failures:
print("\nFailures:")
for image, method, error in failures:
print(f"- {image} / {method}: {error}")
return 1
print("\nAll requested methods completed successfully.")
return 0
if __name__ == "__main__":
raise SystemExit(main())

13
web_dehaze/README.md Normal file
View File

@@ -0,0 +1,13 @@
# Dehaze Web Console
本地网页端入口,用于选择 `待去雾图片/` 中的图片,运行 AOD、Baidu_API、DCP、DehazeNet、GCANet、RefineDNet并统一展示结果与后处理结果。
```bash
bash run_dehaze_web.sh
```
结果统一保存到 `web_results/`。当前环境缺少 `caffe``torch`AOD/DehazeNet/GCANet/RefineDNet 会在运行时显示缺少依赖DCP 和后处理可直接使用当前 Python 环境运行。
当前默认调度:
- AOD、DehazeNet`/home/wkmgc/miniconda3/envs/dehaze_caffe/bin/python`
- GCANet、RefineDNet`/home/wkmgc/miniconda3/envs/seg_server/bin/python`

501
web_dehaze/pipeline.py Normal file
View File

@@ -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

164
web_dehaze/postprocess.py Normal file
View File

@@ -0,0 +1,164 @@
from __future__ import annotations
from pathlib import Path
from typing import Any
import numpy as np
from PIL import Image
from scipy.optimize import minimize
from skimage import color, exposure
POSTPROCESSORS: dict[str, dict[str, Any]] = {
"manual_sv": {
"label": "手动 S/V",
"needs_reference": False,
"params": {"s_gain": 1.0, "v_gain": 1.0},
},
"hsv_hist": {
"label": "HSV 直方图匹配",
"needs_reference": True,
"params": {"match_hue": False},
},
"auto_sv": {
"label": "自动 S/V",
"needs_reference": True,
"params": {},
},
"hist_auto_sv": {
"label": "直方图 + 自动 S/V",
"needs_reference": True,
"params": {"match_hue": False},
},
}
def _load_rgb_float(path: Path) -> tuple[np.ndarray, tuple[int, int]]:
image = Image.open(path).convert("RGB")
return np.asarray(image, dtype=np.float64) / 255.0, image.size
def _load_reference_float(path: Path, size: tuple[int, int]) -> np.ndarray:
image = Image.open(path).convert("RGB")
if image.size != size:
image = image.resize(size, Image.BILINEAR)
return np.asarray(image, dtype=np.float64) / 255.0
def _save_rgb_float(array: np.ndarray, output_path: Path) -> None:
output_path.parent.mkdir(parents=True, exist_ok=True)
image = np.clip(array, 0.0, 1.0)
Image.fromarray((image * 255.0 + 0.5).astype(np.uint8)).save(output_path)
def adjust_sv(source_path: Path, output_path: Path, s_gain: float = 1.0, v_gain: float = 1.0) -> dict[str, Any]:
src_rgb, _ = _load_rgb_float(source_path)
hsv = color.rgb2hsv(src_rgb)
hsv[:, :, 1] = np.clip(hsv[:, :, 1] * float(s_gain), 0.0, 1.0)
hsv[:, :, 2] = np.clip(hsv[:, :, 2] * float(v_gain), 0.0, 1.0)
_save_rgb_float(color.hsv2rgb(hsv), output_path)
return {"s_gain": float(s_gain), "v_gain": float(v_gain)}
def hsv_hist_match(
source_path: Path,
reference_path: Path,
output_path: Path,
match_hue: bool = False,
) -> dict[str, Any]:
src_rgb, size = _load_rgb_float(source_path)
ref_rgb = _load_reference_float(reference_path, size)
hsv_src = color.rgb2hsv(src_rgb)
hsv_ref = color.rgb2hsv(ref_rgb)
hue = exposure.match_histograms(hsv_src[:, :, 0], hsv_ref[:, :, 0]) if match_hue else hsv_src[:, :, 0]
sat = exposure.match_histograms(hsv_src[:, :, 1], hsv_ref[:, :, 1])
val = exposure.match_histograms(hsv_src[:, :, 2], hsv_ref[:, :, 2])
result_hsv = np.stack([hue, sat, val], axis=2)
_save_rgb_float(color.hsv2rgb(result_hsv), output_path)
return {"match_hue": bool(match_hue)}
def auto_sv(
source_path: Path,
reference_path: Path,
output_path: Path,
hist_first: bool = False,
match_hue: bool = False,
) -> dict[str, Any]:
src_rgb, size = _load_rgb_float(source_path)
ref_rgb = _load_reference_float(reference_path, size)
hsv_src = color.rgb2hsv(src_rgb)
hsv_ref = color.rgb2hsv(ref_rgb)
if hist_first:
hue = exposure.match_histograms(hsv_src[:, :, 0], hsv_ref[:, :, 0]) if match_hue else hsv_src[:, :, 0]
sat = exposure.match_histograms(hsv_src[:, :, 1], hsv_ref[:, :, 1])
val = exposure.match_histograms(hsv_src[:, :, 2], hsv_ref[:, :, 2])
hsv_src = np.stack([hue, sat, val], axis=2)
def loss_function(params: np.ndarray) -> float:
ks, kv = params
adj_s = np.clip(hsv_src[:, :, 1] * ks, 0.0, 1.0)
adj_v = np.clip(hsv_src[:, :, 2] * kv, 0.0, 1.0)
loss_s = np.mean((adj_s - hsv_ref[:, :, 1]) ** 2)
loss_v = np.mean((adj_v - hsv_ref[:, :, 2]) ** 2)
return float(loss_s + loss_v)
result = minimize(loss_function, np.array([1.0, 1.0]), method="Nelder-Mead", tol=1e-4)
best_s, best_v = result.x
hsv_final = hsv_src.copy()
hsv_final[:, :, 1] = np.clip(hsv_final[:, :, 1] * best_s, 0.0, 1.0)
hsv_final[:, :, 2] = np.clip(hsv_final[:, :, 2] * best_v, 0.0, 1.0)
_save_rgb_float(color.hsv2rgb(hsv_final), output_path)
return {
"s_gain": float(best_s),
"v_gain": float(best_v),
"hist_first": bool(hist_first),
"match_hue": bool(match_hue),
}
def run_postprocess(
processor: str,
source_path: Path,
output_path: Path,
reference_path: Path | None = None,
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
params = params or {}
if processor == "manual_sv":
return adjust_sv(
source_path,
output_path,
s_gain=float(params.get("s_gain", 1.0)),
v_gain=float(params.get("v_gain", 1.0)),
)
if reference_path is None:
raise ValueError(f"{processor} requires a reference image")
if processor == "hsv_hist":
return hsv_hist_match(
source_path,
reference_path,
output_path,
match_hue=bool(params.get("match_hue", False)),
)
if processor == "auto_sv":
return auto_sv(source_path, reference_path, output_path, hist_first=False)
if processor == "hist_auto_sv":
return auto_sv(
source_path,
reference_path,
output_path,
hist_first=True,
match_hue=bool(params.get("match_hue", False)),
)
raise ValueError(f"Unknown postprocessor: {processor}")

248
web_dehaze/server.py Normal file
View File

@@ -0,0 +1,248 @@
from __future__ import annotations
import argparse
import json
import mimetypes
import threading
import traceback
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import count
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, unquote, urlparse
import pipeline
STATIC_DIR = Path(__file__).resolve().parent / "static"
JOBS: dict[str, dict[str, Any]] = {}
JOB_IDS = count(1)
JOB_LOCK = threading.Lock()
def _json_bytes(payload: Any) -> bytes:
return json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
def _new_job(kind: str, target, *args, **kwargs) -> dict[str, Any]:
job_id = str(next(JOB_IDS))
job = {
"id": job_id,
"kind": kind,
"status": "running",
"logs": [],
"result": None,
"error": None,
}
with JOB_LOCK:
JOBS[job_id] = job
def log(message: str) -> None:
with JOB_LOCK:
job["logs"].append(message)
job["logs"] = job["logs"][-400:]
def wrapped() -> None:
try:
job["result"] = target(log, *args, **kwargs)
job["status"] = "done"
except Exception as exc:
job["status"] = "error"
job["error"] = str(exc)
log(traceback.format_exc())
threading.Thread(target=wrapped, daemon=True).start()
return job
def _run_dehaze_job(log, payload: dict[str, Any]) -> dict[str, Any]:
filename = payload["image"]
methods = payload.get("methods") or []
options = payload.get("options") or {}
postprocessors = payload.get("postprocessors") or []
post_sources = payload.get("post_sources") or []
reference = payload.get("reference") or filename
completed = []
failed = []
for method in methods:
try:
result = pipeline.run_dehaze_method(filename, method, options.get(method, options), log)
completed.append({"method": method, "path": pipeline.relpath(result)})
except Exception as exc:
failed.append({"method": method, "error": str(exc)})
log(f"[{method}] 失败: {exc}")
if postprocessors:
sources = post_sources or [item["method"] for item in completed]
for source in sources:
try:
pipeline.run_postprocessors_for_source(
filename,
source,
postprocessors,
params=payload.get("post_params") or {},
reference_filename=reference,
log=log,
)
except Exception as exc:
failed.append({"postprocess_source": source, "error": str(exc)})
log(f"[后处理/{source}] 失败: {exc}")
return {"completed": completed, "failed": failed, "results": pipeline.get_results(filename)}
def _run_post_job(log, payload: dict[str, Any]) -> dict[str, Any]:
filename = payload["image"]
source = payload.get("source") or "DCP"
processors = payload.get("processors") or []
reference = payload.get("reference") or filename
outputs = pipeline.run_postprocessors_for_source(
filename,
source,
processors,
params=payload.get("params") or {},
reference_filename=reference,
log=log,
)
return {
"outputs": [pipeline.relpath(path) for path in outputs],
"results": pipeline.get_results(filename),
}
class DehazeRequestHandler(BaseHTTPRequestHandler):
server_version = "DehazeWeb/1.0"
def log_message(self, fmt: str, *args) -> None:
print("[%s] %s" % (self.log_date_time_string(), fmt % args))
def _send(self, status: int, body: bytes, content_type: str = "application/json; charset=utf-8") -> None:
self.send_response(status)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def _send_json(self, payload: Any, status: int = HTTPStatus.OK) -> None:
self._send(int(status), _json_bytes(payload))
def _send_error_json(self, status: int, message: str) -> None:
self._send_json({"error": message}, status)
def do_HEAD(self) -> None:
parsed = urlparse(self.path)
target = STATIC_DIR / "index.html" if parsed.path == "/" else STATIC_DIR / parsed.path.removeprefix("/static/")
if parsed.path == "/" or parsed.path.startswith("/static/"):
target = target.resolve()
static_root = STATIC_DIR.resolve()
if static_root not in target.parents and target != static_root:
self.send_response(HTTPStatus.FORBIDDEN)
self.end_headers()
return
if target.exists() and target.is_file():
content_type = mimetypes.guess_type(target.name)[0] or "application/octet-stream"
if target.suffix == ".html":
content_type = "text/html; charset=utf-8"
self.send_response(HTTPStatus.OK)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(target.stat().st_size))
self.end_headers()
return
self.send_response(HTTPStatus.NOT_FOUND)
self.end_headers()
def _read_json(self) -> dict[str, Any]:
length = int(self.headers.get("Content-Length", "0"))
if length <= 0:
return {}
raw = self.rfile.read(length).decode("utf-8")
return json.loads(raw)
def do_GET(self) -> None:
parsed = urlparse(self.path)
path = parsed.path
query = parse_qs(parsed.query)
try:
if path == "/":
self._serve_static("index.html")
elif path.startswith("/static/"):
self._serve_static(path.removeprefix("/static/"))
elif path == "/asset":
rel = unquote(query.get("path", [""])[0])
self._serve_asset(rel)
elif path == "/api/images":
self._send_json({"images": pipeline.list_images()})
elif path == "/api/capabilities":
self._send_json(pipeline.capabilities())
elif path == "/api/results":
filename = query.get("image", [""])[0]
self._send_json(pipeline.get_results(filename))
elif path == "/api/job":
job_id = query.get("id", [""])[0]
with JOB_LOCK:
job = JOBS.get(job_id)
payload = dict(job) if job else None
if payload is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "job not found")
else:
self._send_json(payload)
else:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
except Exception as exc:
self._send_error_json(HTTPStatus.INTERNAL_SERVER_ERROR, str(exc))
def do_POST(self) -> None:
parsed = urlparse(self.path)
try:
payload = self._read_json()
if parsed.path == "/api/run":
job = _new_job("dehaze", _run_dehaze_job, payload)
self._send_json({"job": job})
elif parsed.path == "/api/postprocess":
job = _new_job("postprocess", _run_post_job, payload)
self._send_json({"job": job})
else:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
except Exception as exc:
self._send_error_json(HTTPStatus.BAD_REQUEST, str(exc))
def _serve_static(self, name: str) -> None:
target = (STATIC_DIR / name).resolve()
if STATIC_DIR.resolve() not in target.parents and target != STATIC_DIR.resolve():
self._send_error_json(HTTPStatus.FORBIDDEN, "forbidden")
return
if not target.exists() or not target.is_file():
self._send_error_json(HTTPStatus.NOT_FOUND, "static file not found")
return
content_type = mimetypes.guess_type(target.name)[0] or "application/octet-stream"
if target.suffix == ".html":
content_type = "text/html; charset=utf-8"
elif target.suffix == ".css":
content_type = "text/css; charset=utf-8"
elif target.suffix == ".js":
content_type = "application/javascript; charset=utf-8"
self._send(HTTPStatus.OK, target.read_bytes(), content_type)
def _serve_asset(self, rel: str) -> None:
target = pipeline.resolve_asset(rel)
content_type = mimetypes.guess_type(target.name)[0] or "application/octet-stream"
self._send(HTTPStatus.OK, target.read_bytes(), content_type)
def main() -> None:
parser = argparse.ArgumentParser(description="Local dehaze web console")
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=7860)
args = parser.parse_args()
server = ThreadingHTTPServer((args.host, args.port), DehazeRequestHandler)
print(f"Dehaze web console: http://{args.host}:{args.port}")
print(f"Images: {pipeline.IMAGE_DIR}")
print(f"Results: {pipeline.RESULTS_DIR}")
server.serve_forever()
if __name__ == "__main__":
main()

265
web_dehaze/static/app.js Normal file
View File

@@ -0,0 +1,265 @@
const state = {
images: [],
capabilities: null,
selectedImage: null,
currentJob: null,
pollTimer: null,
};
const $ = (id) => document.getElementById(id);
function assetUrl(path) {
return `/asset?path=${encodeURIComponent(path)}&t=${Date.now()}`;
}
async function api(path, options = {}) {
const response = await fetch(path, options);
const payload = await response.json();
if (!response.ok) {
throw new Error(payload.error || response.statusText);
}
return payload;
}
function formatBytes(bytes) {
if (!Number.isFinite(bytes)) return "";
if (bytes < 1024) return `${bytes} B`;
if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`;
return `${(bytes / 1024 / 1024).toFixed(1)} MB`;
}
function renderImages() {
const box = $("imageList");
box.innerHTML = "";
state.images.forEach((image) => {
const item = document.createElement("button");
item.className = `image-item ${state.selectedImage === image.name ? "active" : ""}`;
item.type = "button";
item.innerHTML = `
<span>${image.name}</span>
<span class="image-meta">${image.width || "?"}x${image.height || "?"} · ${formatBytes(image.size)}</span>
`;
item.addEventListener("click", () => selectImage(image.name));
box.appendChild(item);
});
}
function renderMethods() {
const box = $("methodList");
box.innerHTML = "";
const methods = state.capabilities?.methods || {};
Object.entries(methods).forEach(([key, info]) => {
const label = document.createElement("label");
label.className = `check-item ${info.available ? "" : "disabled"}`;
label.innerHTML = `
<span>
<strong>${info.label}</strong>
<span class="method-meta">${info.available ? "ready" : `${info.module_ok ? "模型" : info.module}`}</span>
</span>
<input type="checkbox" value="${key}" ${info.available || key === "Baidu_API" || key === "DCP" ? "" : "disabled"} ${key === "DCP" ? "checked" : ""} />
`;
box.appendChild(label);
});
}
function renderPostprocessors() {
const box = $("postList");
box.innerHTML = "";
const processors = state.capabilities?.postprocessors || {};
Object.entries(processors).forEach(([key, info]) => {
const label = document.createElement("label");
label.className = "check-item";
label.innerHTML = `
<span><strong>${info.label}</strong></span>
<input type="checkbox" value="${key}" ${key === "manual_sv" ? "checked" : ""} />
`;
box.appendChild(label);
});
}
function renderReferenceOptions() {
const select = $("referenceImage");
select.innerHTML = "";
state.images.forEach((image) => {
const option = document.createElement("option");
option.value = image.name;
option.textContent = image.name;
option.selected = image.name === state.selectedImage;
select.appendChild(option);
});
}
function setJobState(text, status = "") {
const box = $("jobState");
box.textContent = text;
box.className = `job-state ${status}`;
}
async function selectImage(name) {
state.selectedImage = name;
$("currentTitle").textContent = name;
renderImages();
renderReferenceOptions();
await refreshResults();
}
function resultCard(title, path, exists = true) {
const card = document.createElement("article");
card.className = "result-card";
const badge = exists ? '<span class="badge ok">ready</span>' : '<span class="badge pending">未生成</span>';
const body = exists
? `<div class="image-frame"><img src="${assetUrl(path)}" alt="${title}" loading="lazy" /></div>`
: '<div class="image-frame"><span class="missing">暂无结果</span></div>';
card.innerHTML = `<header><h3>${title}</h3>${badge}</header>${body}`;
return card;
}
function updatePostSourceOptions(results) {
const select = $("postSource");
const previous = select.value;
select.innerHTML = "";
const original = document.createElement("option");
original.value = "original";
original.textContent = "原图";
select.appendChild(original);
results.dehaze.filter((item) => item.exists).forEach((item) => {
const option = document.createElement("option");
option.value = item.method;
option.textContent = item.label;
select.appendChild(option);
});
if ([...select.options].some((option) => option.value === previous)) {
select.value = previous;
} else if (results.dehaze.some((item) => item.method === "DCP" && item.exists)) {
select.value = "DCP";
}
}
async function refreshResults() {
if (!state.selectedImage) return;
const results = await api(`/api/results?image=${encodeURIComponent(state.selectedImage)}`);
const grid = $("resultGrid");
grid.innerHTML = "";
grid.appendChild(resultCard(results.original.label, results.original.path, true));
results.dehaze.forEach((item) => {
grid.appendChild(resultCard(item.label, item.path, item.exists));
});
results.postprocess.forEach((item) => {
grid.appendChild(resultCard(item.name, item.path, true));
});
updatePostSourceOptions(results);
}
function selectedValues(containerId) {
return [...$(containerId).querySelectorAll("input[type=checkbox]:checked")].map((input) => input.value);
}
function postParams() {
return {
manual_sv: {
s_gain: Number($("sGain").value),
v_gain: Number($("vGain").value),
},
hsv_hist: {
match_hue: $("matchHue").checked,
},
hist_auto_sv: {
match_hue: $("matchHue").checked,
},
};
}
async function startJob(endpoint, payload) {
const response = await api(endpoint, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
state.currentJob = response.job.id;
$("logBox").textContent = "";
setJobState("Running", "running");
pollJob();
}
async function pollJob() {
if (!state.currentJob) return;
clearTimeout(state.pollTimer);
try {
const job = await api(`/api/job?id=${encodeURIComponent(state.currentJob)}`);
$("logBox").textContent = (job.logs || []).join("\n");
$("logBox").scrollTop = $("logBox").scrollHeight;
if (job.status === "running") {
setJobState("Running", "running");
state.pollTimer = setTimeout(pollJob, 1000);
} else {
setJobState(job.status === "done" ? "Done" : "Error", job.status);
await refreshResults();
}
} catch (error) {
setJobState("Error", "error");
$("logBox").textContent = error.message;
}
}
async function runSelectedModels() {
if (!state.selectedImage) return;
const methods = selectedValues("methodList");
const payload = {
image: state.selectedImage,
methods,
options: {
DCP: {
sz: Number($("dcpSz").value),
tx: Number($("dcpTx").value),
},
},
};
await startJob("/api/run", payload);
}
async function runPostprocess() {
if (!state.selectedImage) return;
const processors = selectedValues("postList");
const payload = {
image: state.selectedImage,
source: $("postSource").value,
reference: $("referenceImage").value || state.selectedImage,
processors,
params: postParams(),
};
await startJob("/api/postprocess", payload);
}
function bindControls() {
$("runBtn").addEventListener("click", runSelectedModels);
$("postBtn").addEventListener("click", runPostprocess);
$("refreshBtn").addEventListener("click", refreshResults);
$("sGain").addEventListener("input", () => {
$("sGainValue").textContent = `${Math.round(Number($("sGain").value) * 100)}%`;
});
$("vGain").addEventListener("input", () => {
$("vGainValue").textContent = `${Math.round(Number($("vGain").value) * 100)}%`;
});
}
async function init() {
bindControls();
const [capabilities, images] = await Promise.all([api("/api/capabilities"), api("/api/images")]);
state.capabilities = capabilities;
state.images = images.images || [];
$("envLine").textContent = capabilities.results_dir;
renderMethods();
renderPostprocessors();
renderImages();
if (state.images.length) {
await selectImage(state.images[0].name);
} else {
$("currentTitle").textContent = "待去雾图片为空";
}
setJobState("Idle");
}
init().catch((error) => {
setJobState("Error", "error");
$("logBox").textContent = error.stack || error.message;
});

View File

@@ -0,0 +1,93 @@
<!doctype html>
<html lang="zh-CN">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1" />
<title>Dehaze Console</title>
<link rel="stylesheet" href="/static/style.css" />
</head>
<body>
<main class="shell">
<aside class="control-panel">
<div class="brand">
<span class="brand-mark"></span>
<div>
<h1>Dehaze Console</h1>
<p id="envLine">Loading...</p>
</div>
</div>
<section class="panel-section">
<h2>待去雾图片</h2>
<div id="imageList" class="image-list"></div>
</section>
<section class="panel-section">
<h2>模型</h2>
<div id="methodList" class="check-grid"></div>
<div class="param-grid">
<label>
<span>DCP 窗口</span>
<input id="dcpSz" type="number" min="3" max="99" step="1" value="10" />
</label>
<label>
<span>DCP tx</span>
<input id="dcpTx" type="number" min="0.01" max="1" step="0.01" value="0.20" />
</label>
</div>
<button id="runBtn" class="primary-btn">运行选中模型</button>
</section>
<section class="panel-section">
<h2>后处理</h2>
<label class="field">
<span>源图</span>
<select id="postSource"></select>
</label>
<label class="field">
<span>参考图</span>
<select id="referenceImage"></select>
</label>
<div id="postList" class="check-grid compact"></div>
<div class="slider-row">
<span>S</span>
<input id="sGain" type="range" min="0" max="2.5" step="0.01" value="1" />
<strong id="sGainValue">100%</strong>
</div>
<div class="slider-row">
<span>V</span>
<input id="vGain" type="range" min="0" max="2.5" step="0.01" value="1" />
<strong id="vGainValue">100%</strong>
</div>
<label class="toggle-line">
<input id="matchHue" type="checkbox" />
<span>匹配 H 通道</span>
</label>
<button id="postBtn" class="secondary-btn">生成后处理</button>
</section>
</aside>
<section class="workspace">
<header class="topbar">
<div>
<p class="eyebrow">当前图片</p>
<h2 id="currentTitle">未选择</h2>
</div>
<div id="jobState" class="job-state">Idle</div>
</header>
<section id="resultGrid" class="result-grid"></section>
<section class="log-panel">
<div class="log-head">
<h2>日志</h2>
<button id="refreshBtn" class="text-btn">刷新</button>
</div>
<pre id="logBox"></pre>
</section>
</section>
</main>
<script src="/static/app.js"></script>
</body>
</html>

428
web_dehaze/static/style.css Normal file
View File

@@ -0,0 +1,428 @@
:root {
--paper: #f4f0e8;
--panel: #fffaf1;
--ink: #1e2520;
--muted: #6d756f;
--line: #d8d0c3;
--green: #1f6b57;
--green-dark: #124839;
--cobalt: #295f9f;
--amber: #c1842d;
--red: #b3473f;
--shadow: 0 18px 60px rgba(33, 37, 31, 0.12);
}
* {
box-sizing: border-box;
}
body {
margin: 0;
min-height: 100vh;
color: var(--ink);
background:
linear-gradient(90deg, rgba(31, 107, 87, 0.06) 1px, transparent 1px),
linear-gradient(0deg, rgba(31, 107, 87, 0.05) 1px, transparent 1px),
var(--paper);
background-size: 28px 28px;
font-family: "Aptos", "Noto Sans SC", "Microsoft YaHei", sans-serif;
}
button,
input,
select {
font: inherit;
}
.shell {
display: grid;
grid-template-columns: minmax(300px, 360px) 1fr;
min-height: 100vh;
}
.control-panel {
padding: 24px 20px;
border-right: 1px solid var(--line);
background: rgba(255, 250, 241, 0.94);
box-shadow: var(--shadow);
overflow-y: auto;
}
.brand {
display: flex;
align-items: center;
gap: 14px;
padding-bottom: 22px;
border-bottom: 1px solid var(--line);
}
.brand-mark {
width: 42px;
height: 42px;
border: 2px solid var(--ink);
background:
linear-gradient(135deg, transparent 46%, var(--ink) 47%, var(--ink) 53%, transparent 54%),
linear-gradient(45deg, var(--green) 0 48%, var(--amber) 48% 100%);
}
h1,
h2,
p {
margin: 0;
}
h1 {
font-size: 24px;
line-height: 1;
}
.brand p,
.eyebrow {
margin-top: 6px;
color: var(--muted);
font-size: 12px;
}
.panel-section {
padding: 20px 0;
border-bottom: 1px solid var(--line);
}
.panel-section h2,
.log-head h2 {
margin-bottom: 12px;
font-size: 14px;
letter-spacing: 0;
}
.image-list,
.check-grid {
display: grid;
gap: 8px;
}
.image-item,
.check-item {
display: flex;
align-items: center;
justify-content: space-between;
gap: 12px;
min-height: 42px;
padding: 9px 10px;
border: 1px solid var(--line);
background: rgba(255, 255, 255, 0.55);
}
.image-item {
cursor: pointer;
}
.image-item.active {
border-color: var(--green);
background: rgba(31, 107, 87, 0.12);
}
.image-meta,
.method-meta {
color: var(--muted);
font-size: 11px;
}
.check-item.disabled {
color: var(--muted);
background: rgba(216, 208, 195, 0.34);
}
.check-item input {
width: 18px;
height: 18px;
accent-color: var(--green);
}
.compact {
grid-template-columns: 1fr;
}
.param-grid {
display: grid;
grid-template-columns: 1fr 1fr;
gap: 8px;
margin-top: 12px;
}
.param-grid label,
.field {
display: grid;
gap: 6px;
color: var(--muted);
font-size: 12px;
}
.param-grid input,
.field select {
width: 100%;
min-height: 38px;
border: 1px solid var(--line);
background: #fffdf8;
color: var(--ink);
padding: 7px 8px;
}
.field {
margin-bottom: 10px;
}
.primary-btn,
.secondary-btn,
.text-btn {
min-height: 42px;
border: 1px solid transparent;
cursor: pointer;
}
.primary-btn,
.secondary-btn {
width: 100%;
margin-top: 12px;
color: #fff;
}
.primary-btn {
background: var(--green);
}
.primary-btn:hover {
background: var(--green-dark);
}
.secondary-btn {
background: var(--cobalt);
}
.secondary-btn:hover {
background: #1e4d83;
}
.text-btn {
padding: 0 12px;
border-color: var(--line);
background: #fffdf8;
color: var(--ink);
}
.slider-row {
display: grid;
grid-template-columns: 20px 1fr 52px;
align-items: center;
gap: 8px;
margin-top: 12px;
color: var(--muted);
font-size: 12px;
}
.slider-row input {
accent-color: var(--green);
}
.slider-row strong {
color: var(--ink);
font-size: 12px;
text-align: right;
}
.toggle-line {
display: flex;
align-items: center;
gap: 8px;
margin-top: 12px;
color: var(--muted);
font-size: 12px;
}
.toggle-line input {
accent-color: var(--green);
}
.workspace {
min-width: 0;
padding: 26px;
overflow: hidden;
}
.topbar {
display: flex;
align-items: end;
justify-content: space-between;
gap: 16px;
margin-bottom: 20px;
}
.topbar h2 {
margin-top: 4px;
font-size: clamp(22px, 3vw, 38px);
line-height: 1.05;
word-break: break-all;
}
.job-state {
min-width: 92px;
padding: 8px 12px;
border: 1px solid var(--line);
background: rgba(255, 250, 241, 0.82);
text-align: center;
color: var(--muted);
}
.job-state.running {
color: var(--cobalt);
border-color: rgba(41, 95, 159, 0.35);
}
.job-state.error {
color: var(--red);
border-color: rgba(179, 71, 63, 0.4);
}
.job-state.done {
color: var(--green);
border-color: rgba(31, 107, 87, 0.4);
}
.result-grid {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(240px, 1fr));
gap: 14px;
max-height: calc(100vh - 300px);
overflow-y: auto;
padding-right: 4px;
}
.result-card {
display: grid;
grid-template-rows: auto 1fr;
min-height: 230px;
border: 1px solid var(--line);
background: rgba(255, 250, 241, 0.82);
}
.result-card header {
display: flex;
align-items: center;
justify-content: space-between;
gap: 10px;
min-height: 42px;
padding: 10px;
border-bottom: 1px solid var(--line);
}
.result-card h3 {
margin: 0;
font-size: 13px;
word-break: break-word;
}
.badge {
flex: 0 0 auto;
padding: 3px 7px;
border: 1px solid var(--line);
color: var(--muted);
font-size: 11px;
}
.badge.ok {
color: var(--green);
border-color: rgba(31, 107, 87, 0.4);
}
.badge.pending {
color: var(--amber);
border-color: rgba(193, 132, 45, 0.42);
}
.image-frame {
display: grid;
place-items: center;
min-height: 188px;
padding: 8px;
background:
linear-gradient(45deg, rgba(30, 37, 32, 0.05) 25%, transparent 25% 75%, rgba(30, 37, 32, 0.05) 75%),
linear-gradient(45deg, rgba(30, 37, 32, 0.05) 25%, transparent 25% 75%, rgba(30, 37, 32, 0.05) 75%);
background-position: 0 0, 8px 8px;
background-size: 16px 16px;
}
.image-frame img {
display: block;
max-width: 100%;
max-height: 38vh;
object-fit: contain;
background: #fff;
}
.missing {
color: var(--muted);
font-size: 13px;
}
.log-panel {
margin-top: 18px;
border: 1px solid var(--line);
background: rgba(30, 37, 32, 0.92);
color: #e9eadf;
}
.log-head {
display: flex;
align-items: center;
justify-content: space-between;
min-height: 44px;
padding: 8px 10px;
border-bottom: 1px solid rgba(255, 255, 255, 0.12);
}
.log-head h2 {
margin: 0;
color: #f7f2e8;
}
.log-panel .text-btn {
min-height: 30px;
background: transparent;
color: #f7f2e8;
border-color: rgba(255, 255, 255, 0.22);
}
#logBox {
height: 172px;
margin: 0;
padding: 12px;
overflow: auto;
white-space: pre-wrap;
font-family: "Cascadia Mono", "Noto Sans Mono", monospace;
font-size: 12px;
line-height: 1.5;
}
@media (max-width: 900px) {
.shell {
grid-template-columns: 1fr;
}
.control-panel {
border-right: 0;
border-bottom: 1px solid var(--line);
}
.workspace {
padding: 20px;
}
.topbar {
align-items: start;
flex-direction: column;
}
.result-grid {
max-height: none;
}
}

View File

@@ -0,0 +1,63 @@
import numpy as np
from PIL import Image
from skimage import exposure
from matplotlib import colors
def match_image_hsv(source_path, reference_path, output_path):
# 1. 读取图片并归一化到 0-1 (matplotlib 的 hsv 转换需要 0-1)
src_rgb = np.array(Image.open(source_path)) / 255.0
ref_rgb = np.array(Image.open(reference_path)) / 255.0
# 2. 将图片从 RGB 转换为 HSV
src_hsv = colors.rgb_to_hsv(src_rgb)
ref_hsv = colors.rgb_to_hsv(ref_rgb)
# 3. 分离 HSV 通道
# src_h: 色调, src_s: 饱和度, src_v: 亮度
s_h, s_s, s_v = src_hsv[:,:,0], src_hsv[:,:,1], src_hsv[:,:,2]
r_h, r_s, r_v = ref_hsv[:,:,0], ref_hsv[:,:,1], ref_hsv[:,:,2]
# 4. 对 V (亮度) 和 S (饱和度) 通道进行直方图匹配
# 我们使用参考图的 V 和 S 分布来调整源图
matched_h = exposure.match_histograms(s_h, r_h)
matched_v = exposure.match_histograms(s_v, r_v)
matched_s = exposure.match_histograms(s_s, r_s)
# 5. 合并通道
# 使用原始的 H (色调),加上匹配后的 S 和 V
# V1 不调整H版本
matched_hsv = np.stack([s_h, matched_s, matched_v], axis=2)
# V2 调整H版本
# matched_hsv = np.stack([matched_h, matched_s, matched_v], axis=2)
# 6. 转换回 RGB 并保存
# 转换回 RGB 后需要 clip 到 0-1 范围,防止数值溢出
matched_rgb = np.clip(colors.hsv_to_rgb(matched_hsv), 0, 1)
# 将 0-1 转换回 0-255 的整数并保存
result_image = Image.fromarray((matched_rgb * 255).astype(np.uint8))
result_image.save(output_path)
print(f"HSV 处理完成,图片已保存至: {output_path}")
def match_image_appearance(source_path, reference_path, output_path):
# 1. 读取图片
# source: 需要调整的图片 (第一张, 偏暗)
# reference: 目标风格图片 (第二张, 正常)
src = np.array(Image.open(source_path))
ref = np.array(Image.open(reference_path))
# 2. 进行直方图匹配
# channel_axis=-1 表示对 RGB 每个通道分别进行匹配
matched = exposure.match_histograms(src, ref, channel_axis=-1)
# 3. 保存结果
result_image = Image.fromarray(matched.astype(np.uint8))
result_image.save(output_path)
print(f"处理完成,图片已保存至: {output_path}")
# 使用示例
source_file = "./去雾图像-北航合作/2025-07-02_084220_VID002.mp4_20251027_001308.661.png"
reference_file = "./去雾图像-北航合作-Result_Baidu/2025-07-02_084220_VID002.mp4_20251027_001308.661.png"
# V1
# match_image_appearance(source_file, reference_file, "adjusted_image_rgb.png")
# V2
match_image_hsv(source_file, reference_file, "adjusted_image_hsv.png")

View File

@@ -0,0 +1,156 @@
import numpy as np
from PIL import Image
from skimage import color
from scipy.optimize import minimize
import os
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
# ================= 配置区域 =================
max_workers = 8 # 【并行】 并行处理的进程数 (建议设为 CPU 核心数,如 4, 8, 16)
# 图片后缀
prefix = "_AOD-Net" # "_10_0.2_result" # "_result"
# 输入文件夹 (带 prefix 后缀的图片)
src_dir = "AOD-Net"
# 参考文件夹 (GT/Ground Truth无 prefix 后缀)
ref_dir = "去雾图像-北航合作-雾图"
# 输出文件夹
out_dir = "AOD-Net+后处理"
# ===========================================
# 尝试导入 tqdm如果没安装则定义一个简单的占位符
try:
from tqdm import tqdm
except ImportError:
def tqdm(iterable, **kwargs):
return iterable
def process_single_image(filename, src_folder, ref_folder, output_folder, prefix=prefix):
"""
处理单张图片的函数,用于并行调用
"""
source_path = os.path.join(src_folder, filename)
try:
# --- 寻找对应的参考图 ---
# 逻辑:去除文件名后缀 "prefix"
name_no_ext, ext = os.path.splitext(filename)
if name_no_ext.endswith(prefix):
ref_name_no_ext = name_no_ext[:-len(prefix)] # 去掉最后n个字符 (prefix)
else:
ref_name_no_ext = name_no_ext
ref_filename = ref_name_no_ext + ext
ref_path = os.path.join(ref_folder, ref_filename)
# 检查参考图是否存在
if not os.path.exists(ref_path):
return f"[跳过] 找不到参考图: {ref_filename} (对应: {filename})"
# --- 读取图片并归一化 ---
img_src_pil = Image.open(source_path).convert('RGB')
img_src = np.array(img_src_pil) / 255.0
img_ref_pil = Image.open(ref_path).convert('RGB')
# 确保参考图尺寸和源图一致
if img_src_pil.size != img_ref_pil.size:
img_ref_pil = img_ref_pil.resize(img_src_pil.size, Image.BILINEAR)
img_ref = np.array(img_ref_pil) / 255.0
# --- 转换到 HSV ---
hsv_src = color.rgb2hsv(img_src)
hsv_ref = color.rgb2hsv(img_ref)
# --- 定义损失函数 ---
# 注意在多进程中loss_function 必须定义在 worker 内部才能访问到 hsv_src/ref
def loss_function(params):
ks, kv = params
adj_s = np.clip(hsv_src[:,:,1] * ks, 0, 1)
adj_v = np.clip(hsv_src[:,:,2] * kv, 0, 1)
loss_s = np.mean((adj_s - hsv_ref[:,:,1])**2)
loss_v = np.mean((adj_v - hsv_ref[:,:,2])**2)
return loss_s + loss_v
# --- 开始优化 ---
res = minimize(loss_function, [1.0, 1.0], method='Nelder-Mead', tol=1e-4)
best_s, best_v = res.x
s_percent = int(best_s * 100)
v_percent = int(best_v * 100)
# --- 应用最佳参数 ---
hsv_new = hsv_src.copy()
hsv_new[:, :, 1] = np.clip(hsv_new[:, :, 1] * best_s, 0, 1)
hsv_new[:, :, 2] = np.clip(hsv_new[:, :, 2] * best_v, 0, 1)
# --- 转回 RGB 并保存 ---
img_result_rgb = color.hsv2rgb(hsv_new)
img_save = Image.fromarray((img_result_rgb * 255).astype(np.uint8))
new_filename = f"{name_no_ext}_S_{s_percent}_V_{v_percent}{ext}"
save_path = os.path.join(output_folder, new_filename)
img_save.save(save_path)
return f"OK: {filename} -> S={s_percent}%, V={v_percent}%"
except Exception as e:
return f"[错误] 处理文件 {filename} 时出错: {str(e)}"
def calculate_and_process_batch_parallel(src_folder, ref_folder, output_folder, max_workers=None):
# 1. 确保输出目录存在
if not os.path.exists(output_folder):
os.makedirs(output_folder)
print(f"已创建输出目录: {output_folder}")
# 2. 获取源文件夹内所有图片文件
valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tif')
file_list = [f for f in os.listdir(src_folder) if f.lower().endswith(valid_extensions)]
total_files = len(file_list)
print(f"共发现 {total_files} 张图片,准备开始并行处理 (进程数: {max_workers if max_workers else '自动'})...\n")
# 3. 并行处理
results = []
# ProcessPoolExecutor 自动管理进程池
# max_workers=None 意味着使用 CPU 核心数
with ProcessPoolExecutor(max_workers=max_workers) as executor:
# 提交所有任务
future_to_file = {
executor.submit(process_single_image, filename, src_folder, ref_folder, output_folder): filename
for filename in file_list
}
# 使用 tqdm 显示进度条as_completed 会在任何一个任务完成时yield
pbar = tqdm(total=total_files, unit="img")
for future in as_completed(future_to_file):
result_msg = future.result()
pbar.update(1)
# 如果是错误或跳过信息打印出来如果是OK只更新进度条不刷屏可选
if not result_msg.startswith("OK"):
tqdm.write(result_msg) # 使用 tqdm.write 防止打断进度条
# else:
# tqdm.write(result_msg) # 如果想看每张图的详细结果,取消注释这行
pbar.close()
print("\n" + "="*30)
print("所有处理已完成。")
# 执行
if __name__ == "__main__":
# Windows 下使用多进程必须放在 if __name__ == "__main__": 之下
if os.path.exists(src_dir) and os.path.exists(ref_dir):
# max_workers 可以手动指定,例如 max_workers=4。如果不填则默认跑满 CPU。
calculate_and_process_batch_parallel(src_dir, ref_dir, out_dir, max_workers = max_workers)
else:
print("错误: 找不到输入文件夹或参考文件夹,请检查路径。")

View File

@@ -0,0 +1,161 @@
import numpy as np
from PIL import Image
from skimage import color, exposure
from scipy.optimize import minimize
import os
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
# ================= 配置区域 =================
# 1. 输入与输出文件夹
SRC_DIR = "去雾图像-北航合作-Result_Baidu" # 待处理图片文件夹
REF_DIR = "去雾图像-北航合作" # 参考图(GT)文件夹
OUT_DIR = "去雾图像-北航合作-Result_Baidu_Own_V3" # 结果输出文件夹
# 2. 功能开关
ENABLE_HIST_MATCH = True # 【开关】 True: 开启直方图匹配; False: 关闭
MAX_WORKERS = 4 # 【并行】 并行处理的进程数 (建议设为 CPU 核心数,如 4, 8, 16)
# ===========================================
def process_single_image(file_info):
"""
单个图片处理函数 (用于并行调用)
file_info: (filename, src_dir, ref_dir, out_dir, enable_hist)
"""
filename, src_folder, ref_folder, output_folder, use_hist = file_info
source_path = os.path.join(src_folder, filename)
# 1. 寻找对应的参考图
# 逻辑:去除文件名后缀 "_result" (例如 "image01_result.png" -> "image01.png")
name_no_ext, ext = os.path.splitext(filename)
if name_no_ext.endswith("_result"):
ref_name_no_ext = name_no_ext[:-7] # 去掉 "_result"
else:
ref_name_no_ext = name_no_ext
ref_filename = ref_name_no_ext + ext
ref_path = os.path.join(ref_folder, ref_filename)
if not os.path.exists(ref_path):
return f"[跳过] 找不到参考图: {filename}"
try:
# 2. 读取图片并归一化 (0-1 float)
img_src_pil = Image.open(source_path).convert('RGB')
img_src = np.array(img_src_pil) / 255.0
img_ref_pil = Image.open(ref_path).convert('RGB')
if img_src_pil.size != img_ref_pil.size:
img_ref_pil = img_ref_pil.resize(img_src_pil.size, Image.BILINEAR)
img_ref = np.array(img_ref_pil) / 255.0
# 3. RGB -> HSV
hsv_src = color.rgb2hsv(img_src)
hsv_ref = color.rgb2hsv(img_ref)
# === 新增功能: 直方图匹配 (Histogram Matching) ===
if use_hist:
# 分离通道
s_h, s_s, s_v = hsv_src[:,:,0], hsv_src[:,:,1], hsv_src[:,:,2]
r_h, r_s, r_v = hsv_ref[:,:,0], hsv_ref[:,:,1], hsv_ref[:,:,2]
# 对 S 和 V 通道进行直方图匹配
# 这会将 src 的分布形状强行调整为 ref 的分布形状
matched_s = exposure.match_histograms(s_s, r_s)
matched_v = exposure.match_histograms(s_v, r_v)
# 更新 hsv_src后续的 minimize 将在此基础上进一步微调系数
hsv_src = np.stack([s_h, matched_s, matched_v], axis=-1)
# 4. 优化 S/V 乘数因子
# 即使做了直方图匹配,我们依然计算一个最佳的整体缩放系数,以确保整体误差最小
def loss_function(params):
ks, kv = params
adj_s = np.clip(hsv_src[:,:,1] * ks, 0, 1)
adj_v = np.clip(hsv_src[:,:,2] * kv, 0, 1)
loss_s = np.mean((adj_s - hsv_ref[:,:,1])**2)
loss_v = np.mean((adj_v - hsv_ref[:,:,2])**2)
return loss_s + loss_v
# 初始猜测 [1.0, 1.0]
res = minimize(loss_function, [1.0, 1.0], method='Nelder-Mead', tol=1e-4)
best_s, best_v = res.x
s_percent = int(best_s * 100)
v_percent = int(best_v * 100)
# 5. 应用最终参数
hsv_final = hsv_src.copy()
hsv_final[:, :, 1] = np.clip(hsv_final[:, :, 1] * best_s, 0, 1)
hsv_final[:, :, 2] = np.clip(hsv_final[:, :, 2] * best_v, 0, 1)
# 6. 保存结果
img_result_rgb = color.hsv2rgb(hsv_final)
img_save = Image.fromarray((img_result_rgb * 255).astype(np.uint8))
# 命名增加标识,如果开启了直方图匹配,可以在文件名加个标记(可选)
# 这里保持您要求的格式: 原文件名_S_XX_V_XX.png
new_filename = f"{name_no_ext}_S_{s_percent}_V_{v_percent}{ext}"
save_path = os.path.join(output_folder, new_filename)
img_save.save(save_path)
match_tag = "[HistMatch]" if use_hist else "[Raw]"
return f"{match_tag} 完成: {new_filename} (S={s_percent}%, V={v_percent}%)"
except Exception as e:
return f"[错误] {filename}: {str(e)}"
def main():
# 1. 检查文件夹
if not os.path.exists(SRC_DIR) or not os.path.exists(REF_DIR):
print("错误: 输入或参考文件夹不存在。")
return
if not os.path.exists(OUT_DIR):
os.makedirs(OUT_DIR)
# 2. 获取文件列表
valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tif')
file_list = [f for f in os.listdir(SRC_DIR) if f.lower().endswith(valid_extensions)]
total_files = len(file_list)
if total_files == 0:
print("源文件夹为空。")
return
print(f"=== 开始处理 ===")
print(f"模式: {'直方图匹配 + 参数优化' if ENABLE_HIST_MATCH else '仅参数优化'}")
print(f"并行: {MAX_WORKERS} 线程")
print(f"数量: {total_files} 张图片")
print("-" * 30)
# 3. 准备任务参数
tasks = []
for f in file_list:
# 打包参数传给 worker
tasks.append((f, SRC_DIR, REF_DIR, OUT_DIR, ENABLE_HIST_MATCH))
# 4. 并行执行
start_time = time.time()
with ProcessPoolExecutor(max_workers=MAX_WORKERS) as executor:
# 提交所有任务
futures = [executor.submit(process_single_image, task) for task in tasks]
# 获取结果 (as_completed 会在任务完成时立即返回)
for i, future in enumerate(as_completed(futures)):
result = future.result()
print(f"[{i+1}/{total_files}] {result}")
end_time = time.time()
print("-" * 30)
print(f"全部完成! 耗时: {end_time - start_time:.2f}")
print(f"结果保存在: {OUT_DIR}")
if __name__ == "__main__":
# Windows 下使用多进程必须放在 if __name__ == "__main__": 之下
main()

112
使用手册.md Normal file
View File

@@ -0,0 +1,112 @@
# Dehaze 使用手册
## 1. 项目用途
本项目用于对 `待去雾图片/` 中的图片进行去雾,并在网页端统一展示以下方法的结果:
- AOD
- Baidu_API
- DCP
- DehazeNet
- GCANet
- RefineDNet
- 后处理:手动 S/V、HSV 直方图匹配、自动 S/V、直方图匹配 + 自动 S/V
## 2. 启动网页
```bash
cd /home/wkmgc/Desktop/Dehaze
./run_dehaze_web.sh
```
浏览器访问:
```text
http://192.168.3.11:7860/
```
## 3. 环境说明
当前项目使用统一网页服务调度不同模型环境:
- Web、DCP、Baidu_API、后处理当前 base Python
- AOD、DehazeNet`/home/wkmgc/miniconda3/envs/dehaze_caffe/bin/python`
- GCANet、RefineDNet`/home/wkmgc/miniconda3/envs/seg_server/bin/python`
`run_dehaze_web.sh` 已默认配置这些解释器。如需换环境,可设置:
```bash
export DEHAZE_CAFFE_PYTHON=/path/to/caffe/python
export DEHAZE_TORCH_PYTHON=/path/to/torch/python
```
## 4. 百度 API 配置
真实密钥不要提交到 git。复制示例文件
```bash
cp config/baidu_api.env.example config/baidu_api.env
```
填入:
```text
BAIDU_API_KEY=...
BAIDU_SECRET_KEY=...
```
也可以在启动前直接设置环境变量。
## 5. 使用流程
1. 将待处理图片放入 `待去雾图片/`
2. 启动网页。
3. 在左侧选择图片。
4. 勾选需要运行的模型,点击“运行选中模型”。
5. 在后处理区域选择源图、参考图和后处理方法,点击“生成后处理”。
6. 结果会保存到 `web_results/`,网页会自动刷新显示。
页面中“未生成”表示对应结果文件还不存在,并不是任务卡住。
## 6. 命令行验证
验证全部图片和全部模型:
```bash
python scripts/verify_all.py
```
如不想调用百度 API
```bash
python scripts/verify_all.py --skip-baidu
```
只验证指定图片:
```bash
python scripts/verify_all.py --images 1.png
```
## 7. 清理生成物
清理缓存和运行结果:
```bash
./scripts/clean_generated.sh
```
清理后重新启动网页并运行模型即可再生成结果。
## 8. 目录说明
- `web_dehaze/`:统一网页服务、模型调度和后处理代码。
- `待去雾图片/`:待处理原图。
- `web_results/`:网页运行生成结果,已加入 `.gitignore`
- `AOD-Net_最好加入后处理/`AOD 模型与入口脚本。
- `Baidu_API_最好加入后处理/`:百度去雾 API 脚本。
- `DCP_最好加入后处理/`DCP 去雾脚本。
- `DehazeNet/`DehazeNet 模型与入口脚本。
- `GCANet/`GCANet 模型与入口脚本。
- `RefineDNet/`RefineDNet 模型、权重与入口脚本。
- `※程序-后处理汇总/`:原始后处理脚本归档。

BIN
待去雾图片/1.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 31 KiB

Some files were not shown because too many files have changed in this diff Show More