commit 6db15ebc3f2fe6934917e421cd4aeb2694cc56f6 Author: admin Date: Wed Jun 10 17:42:11 2026 +0800 整合去雾网页工具 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..872f1e1 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/AOD-Net_最好加入后处理/AOD-Net with PONO/README.md b/AOD-Net_最好加入后处理/AOD-Net with PONO/README.md new file mode 100644 index 0000000..5262fd9 --- /dev/null +++ b/AOD-Net_最好加入后处理/AOD-Net with PONO/README.md @@ -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} +} +``` diff --git a/AOD-Net_最好加入后处理/AOD-Net with PONO/model.py b/AOD-Net_最好加入后处理/AOD-Net with PONO/model.py new file mode 100644 index 0000000..f094dfe --- /dev/null +++ b/AOD-Net_最好加入后处理/AOD-Net with PONO/model.py @@ -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 diff --git a/AOD-Net_最好加入后处理/AOD-Net with PONO/pono_train.py b/AOD-Net_最好加入后处理/AOD-Net with PONO/pono_train.py new file mode 100644 index 0000000..5d9830c --- /dev/null +++ b/AOD-Net_最好加入后处理/AOD-Net with PONO/pono_train.py @@ -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) diff --git a/AOD-Net_最好加入后处理/AOD-Net with PONO/ponomodels.zip b/AOD-Net_最好加入后处理/AOD-Net with PONO/ponomodels.zip new file mode 100644 index 0000000..49e47ca Binary files /dev/null and b/AOD-Net_最好加入后处理/AOD-Net with PONO/ponomodels.zip differ diff --git a/AOD-Net_最好加入后处理/AOD-Net with PONO/run_pono_train.sh b/AOD-Net_最好加入后处理/AOD-Net with PONO/run_pono_train.sh new file mode 100644 index 0000000..4bf5265 --- /dev/null +++ b/AOD-Net_最好加入后处理/AOD-Net with PONO/run_pono_train.sh @@ -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 diff --git a/AOD-Net_最好加入后处理/AOD_Net.caffemodel b/AOD-Net_最好加入后处理/AOD_Net.caffemodel new file mode 100644 index 0000000..e8db439 Binary files /dev/null and b/AOD-Net_最好加入后处理/AOD_Net.caffemodel differ diff --git a/AOD-Net_最好加入后处理/All_in_One.sh b/AOD-Net_最好加入后处理/All_in_One.sh new file mode 100644 index 0000000..0fcd247 --- /dev/null +++ b/AOD-Net_最好加入后处理/All_in_One.sh @@ -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 \ No newline at end of file diff --git a/AOD-Net_最好加入后处理/Readme/AOD-Net_result.png b/AOD-Net_最好加入后处理/Readme/AOD-Net_result.png new file mode 100644 index 0000000..3e50d06 Binary files /dev/null and b/AOD-Net_最好加入后处理/Readme/AOD-Net_result.png differ diff --git a/AOD-Net_最好加入后处理/Readme/README.md b/AOD-Net_最好加入后处理/Readme/README.md new file mode 100644 index 0000000..2a0c166 --- /dev/null +++ b/AOD-Net_最好加入后处理/Readme/README.md @@ -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} +} +``` diff --git a/AOD-Net_最好加入后处理/test/test.py b/AOD-Net_最好加入后处理/test/test.py new file mode 100644 index 0000000..7a9e5a2 --- /dev/null +++ b/AOD-Net_最好加入后处理/test/test.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import caffe +import cv2 + + +BASE_DIR = Path(__file__).resolve().parents[1] +SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"} + + +def edit_fcn_proto(template_file: Path, output_file: Path, height: int, width: int) -> None: + template = template_file.read_text() + output_file.write_text(template.format(height=height, width=width)) + + +def iter_images(input_dir: Path) -> list[Path]: + return [ + path + for path in sorted(input_dir.iterdir(), key=lambda p: p.name.lower()) + if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS + ] + + +def run(input_dir: Path, output_dir: Path, max_dim: int) -> None: + caffe.set_mode_cpu() + output_dir.mkdir(parents=True, exist_ok=True) + + template_file = BASE_DIR / "test" / "test_template.prototxt" + deploy_file = BASE_DIR / "test" / "DeployT.prototxt" + model_file = BASE_DIR / "AOD_Net.caffemodel" + + images = iter_images(input_dir) + print(f"image numbers: {len(images)}") + + for index, img_path in enumerate(images, start=1): + npstore = caffe.io.load_image(str(img_path)) + orig_h, orig_w = npstore.shape[0], npstore.shape[1] + + if orig_h > max_dim or orig_w > max_dim: + scale = max_dim / float(max(orig_h, orig_w)) + new_h = int(orig_h * scale) + new_w = int(orig_w * scale) + npstore = cv2.resize(npstore, (new_w, new_h), interpolation=cv2.INTER_CUBIC) + print(f"Resized {img_path.name} from {orig_w}x{orig_h} to {new_w}x{new_h}") + + height, width = npstore.shape[0], npstore.shape[1] + edit_fcn_proto(template_file, deploy_file, height, width) + + net = caffe.Net(str(deploy_file), str(model_file), caffe.TEST) + data = npstore.transpose((2, 0, 1)) + net.blobs["data"].data[...] = [data] + net.forward() + + result = net.blobs["sum"].data[0].transpose((1, 2, 0)) + result = result[:, :, ::-1] + + if height != orig_h or width != orig_w: + result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_CUBIC) + + save_path = output_dir / f"{img_path.stem}_AOD-Net.png" + cv2.imwrite(str(save_path), result * 255.0, [cv2.IMWRITE_JPEG_QUALITY, 100]) + print(f"[{index}/{len(images)}] saved: {save_path}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run AOD-Net on a folder of images.") + parser.add_argument("input_dir", nargs="?", default=str(BASE_DIR / "data" / "img")) + parser.add_argument("output_dir", nargs="?", default=str(BASE_DIR / "data" / "result")) + parser.add_argument("--max-dim", type=int, default=1920) + args = parser.parse_args() + + run(Path(args.input_dir), Path(args.output_dir), args.max_dim) + + +if __name__ == "__main__": + main() diff --git a/AOD-Net_最好加入后处理/test/test_template.prototxt b/AOD-Net_最好加入后处理/test/test_template.prototxt new file mode 100644 index 0000000..98f01fc --- /dev/null +++ b/AOD-Net_最好加入后处理/test/test_template.prototxt @@ -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" +}} + + diff --git a/Baidu_API_最好加入后处理/1_Baidu_Dehaze.py b/Baidu_API_最好加入后处理/1_Baidu_Dehaze.py new file mode 100644 index 0000000..42c0abd --- /dev/null +++ b/Baidu_API_最好加入后处理/1_Baidu_Dehaze.py @@ -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(): + """ + 使用 AK，SK 生成鉴权签名（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() diff --git a/Baidu_API_最好加入后处理/使用方法.txt b/Baidu_API_最好加入后处理/使用方法.txt new file mode 100644 index 0000000..8a74876 --- /dev/null +++ b/Baidu_API_最好加入后处理/使用方法.txt @@ -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 \ No newline at end of file diff --git a/DCP_最好加入后处理/All_in_One.sh b/DCP_最好加入后处理/All_in_One.sh new file mode 100644 index 0000000..4df0d8b --- /dev/null +++ b/DCP_最好加入后处理/All_in_One.sh @@ -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 \ No newline at end of file diff --git a/DCP_最好加入后处理/LICENSE b/DCP_最好加入后处理/LICENSE new file mode 100644 index 0000000..df81365 --- /dev/null +++ b/DCP_最好加入后处理/LICENSE @@ -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. diff --git a/DCP_最好加入后处理/README.md b/DCP_最好加入后处理/README.md new file mode 100644 index 0000000..c9a476d --- /dev/null +++ b/DCP_最好加入后处理/README.md @@ -0,0 +1,18 @@ +# single image dehaze +## Introduction +This program implement single image dehazing using dark channel prior. + +## Compile Dependencies +- OpenCV +- Numpy + +## Examples +
+图片名称 +图片名称 +
+ + +## 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. diff --git a/DCP_最好加入后处理/dehaze.py b/DCP_最好加入后处理/dehaze.py new file mode 100644 index 0000000..44b14ba --- /dev/null +++ b/DCP_最好加入后处理/dehaze.py @@ -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); + diff --git a/DehazeNet/All_in_One.sh b/DehazeNet/All_in_One.sh new file mode 100644 index 0000000..812378d --- /dev/null +++ b/DehazeNet/All_in_One.sh @@ -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 \ No newline at end of file diff --git a/DehazeNet/DehazeFcnTemplate.prototxt b/DehazeNet/DehazeFcnTemplate.prototxt new file mode 100644 index 0000000..ee937de --- /dev/null +++ b/DehazeNet/DehazeFcnTemplate.prototxt @@ -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" +}} diff --git a/DehazeNet/DehazeNet.caffemodel b/DehazeNet/DehazeNet.caffemodel new file mode 100644 index 0000000..5613214 Binary files /dev/null and b/DehazeNet/DehazeNet.caffemodel differ diff --git a/DehazeNet/DehazeNet.prototxt b/DehazeNet/DehazeNet.prototxt new file mode 100644 index 0000000..c668ff2 --- /dev/null +++ b/DehazeNet/DehazeNet.prototxt @@ -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" +} \ No newline at end of file diff --git a/DehazeNet/DehazeNet.py b/DehazeNet/DehazeNet.py new file mode 100644 index 0000000..972f8b0 --- /dev/null +++ b/DehazeNet/DehazeNet.py @@ -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); \ No newline at end of file diff --git a/DehazeNet/readme.md b/DehazeNet/readme.md new file mode 100644 index 0000000..fe0c1b8 --- /dev/null +++ b/DehazeNet/readme.md @@ -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} + } diff --git a/GCANet/All_in_One.sh b/GCANet/All_in_One.sh new file mode 100644 index 0000000..bc2d4bd --- /dev/null +++ b/GCANet/All_in_One.sh @@ -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 \ No newline at end of file diff --git a/GCANet/GCANet.py b/GCANet/GCANet.py new file mode 100644 index 0000000..ca8a826 --- /dev/null +++ b/GCANet/GCANet.py @@ -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 diff --git a/GCANet/GCANet_train/GCANet.py b/GCANet/GCANet_train/GCANet.py new file mode 100644 index 0000000..ca8a826 --- /dev/null +++ b/GCANet/GCANet_train/GCANet.py @@ -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 diff --git a/GCANet/GCANet_train/ImagePairPrefixFolder.py b/GCANet/GCANet_train/ImagePairPrefixFolder.py new file mode 100644 index 0000000..b8d0000 --- /dev/null +++ b/GCANet/GCANet_train/ImagePairPrefixFolder.py @@ -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 diff --git a/GCANet/GCANet_train/folder_loader.py b/GCANet/GCANet_train/folder_loader.py new file mode 100644 index 0000000..a107d62 --- /dev/null +++ b/GCANet/GCANet_train/folder_loader.py @@ -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) diff --git a/GCANet/GCANet_train/test.py b/GCANet/GCANet_train/test.py new file mode 100644 index 0000000..004335d --- /dev/null +++ b/GCANet/GCANet_train/test.py @@ -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)) + + + + + diff --git a/GCANet/GCANet_train/tf_visualizer.py b/GCANet/GCANet_train/tf_visualizer.py new file mode 100644 index 0000000..b3e8b5c --- /dev/null +++ b/GCANet/GCANet_train/tf_visualizer.py @@ -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) + diff --git a/GCANet/GCANet_train/train.py b/GCANet/GCANet_train/train.py new file mode 100644 index 0000000..1495d15 --- /dev/null +++ b/GCANet/GCANet_train/train.py @@ -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)) + + + + + + diff --git a/GCANet/GCANet_train/utils.py b/GCANet/GCANet_train/utils.py new file mode 100644 index 0000000..a1bf877 --- /dev/null +++ b/GCANet/GCANet_train/utils.py @@ -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 \ No newline at end of file diff --git a/GCANet/README.md b/GCANet/README.md new file mode 100644 index 0000000..6dabf5f --- /dev/null +++ b/GCANet/README.md @@ -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](), [Mingming He](), [Qingnan Fan](), *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 + +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 + + diff --git a/GCANet/models/wacv_gcanet_dehaze.pth b/GCANet/models/wacv_gcanet_dehaze.pth new file mode 100644 index 0000000..10ddc06 Binary files /dev/null and b/GCANet/models/wacv_gcanet_dehaze.pth differ diff --git a/GCANet/models/wacv_gcanet_derain.pth b/GCANet/models/wacv_gcanet_derain.pth new file mode 100644 index 0000000..f6269a7 Binary files /dev/null and b/GCANet/models/wacv_gcanet_derain.pth differ diff --git a/GCANet/test.py b/GCANet/test.py new file mode 100644 index 0000000..fddf91b --- /dev/null +++ b/GCANet/test.py @@ -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)) + + + + + diff --git a/GCANet/utils.py b/GCANet/utils.py new file mode 100644 index 0000000..a566416 --- /dev/null +++ b/GCANet/utils.py @@ -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 \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..feed32c --- /dev/null +++ b/README.md @@ -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` 的六个模型和四种后处理流程均可运行。 diff --git a/RefineDNet/All_in_One.sh b/RefineDNet/All_in_One.sh new file mode 100644 index 0000000..5814c08 --- /dev/null +++ b/RefineDNet/All_in_One.sh @@ -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 \ No newline at end of file diff --git a/RefineDNet/README.md b/RefineDNet/README.md new file mode 100644 index 0000000..cb20be0 --- /dev/null +++ b/RefineDNet/README.md @@ -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 +``` + +├── 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 . +``` +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 `/datatsets/quick_test/refined_DCP_outdoor_ep_60`. + +## Test on BeDDE +1. Download the pretrained model on BeDDE. + +2. Run the following command from ``. +``` +python test_BeDDE.py --dataroot --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 `//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 /datasets. +Your directory tree should look like +``` + +├── datasets +│ ├── BeDDE +│ ├── RESIDE-unpaired +│ │ ├── trainA +│ │ └── trainB +│ ... +... +``` +2. Open visdom by `python -m visdom.server` + +3. Run the following command from ``. +``` +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 `/datasets/ITS`. + +2. Rename the hazy image folder as `trainA` and the clear image folder as `trainB`. +Then, your directory tree should look like +``` + +├── datasets +│ ├── BeDDE +│ ├── ITS +│ │ ├── trainA +│ │ └── trainB +│ ... +... +``` +3. Open visdom by `python -m visdom.server` + +4. Run the following command from ``. +``` +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) diff --git a/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_D.pth b/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_D.pth new file mode 100644 index 0000000..039d3d3 Binary files /dev/null and b/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_D.pth differ diff --git a/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_Refiner_J.pth b/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_Refiner_J.pth new file mode 100644 index 0000000..00d973a Binary files /dev/null and b/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_Refiner_J.pth differ diff --git a/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_Refiner_T.pth b/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_Refiner_T.pth new file mode 100644 index 0000000..583feb0 Binary files /dev/null and b/RefineDNet/checkpoints/refined_DCP_outdoor/60_net_Refiner_T.pth differ diff --git a/RefineDNet/checkpoints/refined_DCP_outdoor/src_opt.txt b/RefineDNet/checkpoints/refined_DCP_outdoor/src_opt.txt new file mode 100644 index 0000000..5fe8c2d --- /dev/null +++ b/RefineDNet/checkpoints/refined_DCP_outdoor/src_opt.txt @@ -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 ------------------- diff --git a/RefineDNet/checkpoints/refined_DCP_outdoor/test_opt.txt b/RefineDNet/checkpoints/refined_DCP_outdoor/test_opt.txt new file mode 100644 index 0000000..e624964 --- /dev/null +++ b/RefineDNet/checkpoints/refined_DCP_outdoor/test_opt.txt @@ -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 ------------------- diff --git a/RefineDNet/data/__init__.py b/RefineDNet/data/__init__.py new file mode 100644 index 0000000..8cb6186 --- /dev/null +++ b/RefineDNet/data/__init__.py @@ -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. + -- : (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 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 diff --git a/RefineDNet/data/aligned_dataset.py b/RefineDNet/data/aligned_dataset.py new file mode 100644 index 0000000..cce2be3 --- /dev/null +++ b/RefineDNet/data/aligned_dataset.py @@ -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) diff --git a/RefineDNet/data/base_dataset.py b/RefineDNet/data/base_dataset.py new file mode 100644 index 0000000..e38f27b --- /dev/null +++ b/RefineDNet/data/base_dataset.py @@ -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. + -- : (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 diff --git a/RefineDNet/data/colorization_dataset.py b/RefineDNet/data/colorization_dataset.py new file mode 100644 index 0000000..2616c61 --- /dev/null +++ b/RefineDNet/data/colorization_dataset.py @@ -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) diff --git a/RefineDNet/data/image_folder.py b/RefineDNet/data/image_folder.py new file mode 100644 index 0000000..a9cea74 --- /dev/null +++ b/RefineDNet/data/image_folder.py @@ -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) diff --git a/RefineDNet/data/paired_dataset.py b/RefineDNet/data/paired_dataset.py new file mode 100644 index 0000000..41016ab --- /dev/null +++ b/RefineDNet/data/paired_dataset.py @@ -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 \ No newline at end of file diff --git a/RefineDNet/data/simple_bedde_dataset.py b/RefineDNet/data/simple_bedde_dataset.py new file mode 100644 index 0000000..c7b4b35 --- /dev/null +++ b/RefineDNet/data/simple_bedde_dataset.py @@ -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 diff --git a/RefineDNet/data/simple_paired_dataset.py b/RefineDNet/data/simple_paired_dataset.py new file mode 100644 index 0000000..c7c223e --- /dev/null +++ b/RefineDNet/data/simple_paired_dataset.py @@ -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 \ No newline at end of file diff --git a/RefineDNet/data/single_dataset.py b/RefineDNet/data/single_dataset.py new file mode 100644 index 0000000..a1dc875 --- /dev/null +++ b/RefineDNet/data/single_dataset.py @@ -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) diff --git a/RefineDNet/data/template_dataset.py b/RefineDNet/data/template_dataset.py new file mode 100644 index 0000000..bfdf16b --- /dev/null +++ b/RefineDNet/data/template_dataset.py @@ -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.py +The class name should be Dataset.py +You need to implement the following functions: + -- :　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 ; 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) diff --git a/RefineDNet/data/unaligned_dataset.py b/RefineDNet/data/unaligned_dataset.py new file mode 100644 index 0000000..dc19e26 --- /dev/null +++ b/RefineDNet/data/unaligned_dataset.py @@ -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) diff --git a/RefineDNet/data/unpaired_dataset.py b/RefineDNet/data/unpaired_dataset.py new file mode 100644 index 0000000..c371687 --- /dev/null +++ b/RefineDNet/data/unpaired_dataset.py @@ -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) diff --git a/RefineDNet/datasets/figures/framework_github.jpg b/RefineDNet/datasets/figures/framework_github.jpg new file mode 100644 index 0000000..b559263 Binary files /dev/null and b/RefineDNet/datasets/figures/framework_github.jpg differ diff --git a/RefineDNet/datasets/figures/framework_github2.jpg b/RefineDNet/datasets/figures/framework_github2.jpg new file mode 100644 index 0000000..b559263 Binary files /dev/null and b/RefineDNet/datasets/figures/framework_github2.jpg differ diff --git a/RefineDNet/datasets/figures/outdoor_com_github.jpg b/RefineDNet/datasets/figures/outdoor_com_github.jpg new file mode 100644 index 0000000..37a941b Binary files /dev/null and b/RefineDNet/datasets/figures/outdoor_com_github.jpg differ diff --git a/RefineDNet/models/__init__.py b/RefineDNet/models/__init__.py new file mode 100644 index 0000000..fc01113 --- /dev/null +++ b/RefineDNet/models/__init__.py @@ -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). + -- : unpack data from dataset and apply preprocessing. + -- : produce intermediate results. + -- : calculate loss, gradients, and update network weights. + -- : (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 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 diff --git a/RefineDNet/models/base_model.py b/RefineDNet/models/base_model.py new file mode 100644 index 0000000..307c2ba --- /dev/null +++ b/RefineDNet/models/base_model.py @@ -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). + -- : unpack data from dataset and apply preprocessing. + -- : produce intermediate results. + -- : calculate losses, gradients, and update network weights. + -- : (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 + 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 and .""" + 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 function in no_grad() so we don't save intermediate steps for backprop + It also calls 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 diff --git a/RefineDNet/models/basic_dehaze_model.py b/RefineDNet/models/basic_dehaze_model.py new file mode 100644 index 0000000..441e970 --- /dev/null +++ b/RefineDNet/models/basic_dehaze_model.py @@ -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 + 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 + 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 and . + 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 . + 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 and .""" + # 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 function in no_grad() so we don't save intermediate steps for backprop + It also calls 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 diff --git a/RefineDNet/models/colorization_model.py b/RefineDNet/models/colorization_model.py new file mode 100644 index 0000000..2b4a127 --- /dev/null +++ b/RefineDNet/models/colorization_model.py @@ -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) diff --git a/RefineDNet/models/cycle_gan_model.py b/RefineDNet/models/cycle_gan_model.py new file mode 100644 index 0000000..836cc8f --- /dev/null +++ b/RefineDNet/models/cycle_gan_model.py @@ -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 + 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 + 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 and . + 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 . + 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 and .""" + 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 function in no_grad() so we don't save intermediate steps for backprop + It also calls 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 diff --git a/RefineDNet/models/networks.py b/RefineDNet/models/networks.py new file mode 100644 index 0000000..7d1053b --- /dev/null +++ b/RefineDNet/models/networks.py @@ -0,0 +1,1219 @@ +import torch +import torch.nn as nn +from torch.nn import init +import functools +from torch.optim import lr_scheduler + +import torch.nn.functional as F + +import numpy as np + + +############################################################################### +# Helper Functions +############################################################################### + + +# class Identity(nn.Module): +# def forward(self, x): +# return x + + +def get_norm_layer(norm_type='instance'): + """Return a normalization layer + + Parameters: + norm_type (str) -- the name of the normalization layer: batch | instance | none + + For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). + For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. + """ + if norm_type == 'batch': + norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True) + elif norm_type == 'instance': + norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) + elif norm_type == 'none': + norm_layer = lambda x: Identity() + else: + raise NotImplementedError('normalization layer [%s] is not found' % norm_type) + return norm_layer + + +def get_scheduler(optimizer, opt): + """Return a learning rate scheduler + + Parameters: + optimizer -- the optimizer of the network + opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　 + opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine + + For 'linear', we keep the same learning rate for the first epochs + and linearly decay the rate to zero over the next epochs. + For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. + See https://pytorch.org/docs/stable/optim.html for more details. + """ + if opt.lr_policy == 'linear': + def lambda_rule(epoch): + lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1) + return lr_l + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) + elif opt.lr_policy == 'step': + scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1) + elif opt.lr_policy == 'plateau': + scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) + elif opt.lr_policy == 'cosine': + scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0) + else: + return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy) + return scheduler + + +def init_weights(net, init_type='normal', init_gain=0.02): + """Initialize network weights. + + Parameters: + net (network) -- network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + + We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might + work better for some applications. Feel free to try yourself. + """ + def init_func(m): # define the initialization function + classname = m.__class__.__name__ + if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): + if init_type == 'normal': + init.normal_(m.weight.data, 0.0, init_gain) + elif init_type == 'xavier': + init.xavier_normal_(m.weight.data, gain=init_gain) + elif init_type == 'kaiming': + init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + elif init_type == 'orthogonal': + init.orthogonal_(m.weight.data, gain=init_gain) + else: + raise NotImplementedError('initialization method [%s] is not implemented' % init_type) + if hasattr(m, 'bias') and m.bias is not None: + init.constant_(m.bias.data, 0.0) + elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. + init.normal_(m.weight.data, 1.0, init_gain) + init.constant_(m.bias.data, 0.0) + + print('initialize network with %s' % init_type) + net.apply(init_func) # apply the initialization function + + +def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]): + """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights + Parameters: + net (network) -- the network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Return an initialized network. + """ + if len(gpu_ids) > 0: + assert(torch.cuda.is_available()) + net.to(gpu_ids[0]) + net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs + init_weights(net, init_type, init_gain=init_gain) + return net + + +def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]): + """Create a generator + + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + ngf (int) -- the number of filters in the last conv layer + netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128 + norm (str) -- the name of normalization layers used in the network: batch | instance | none + use_dropout (bool) -- if use dropout layers. + init_type (str) -- the name of our initialization method. + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Returns a generator + + Our current implementation provides two types of generators: + U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images) + The original U-Net paper: https://arxiv.org/abs/1505.04597 + + Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks) + Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations. + We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style). + + + The generator has been initialized by . It uses RELU for non-linearity. + """ + net = None + norm_layer = get_norm_layer(norm_type=norm) + + if netG == 'resnet_9blocks': + net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9) + elif netG == 'resnet_6blocks': + net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6) + elif netG == 'resnet_9blocks_inter': + net = ResnetGWithIntermediate(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9) + elif netG == 'resnet_9blocks_dehaze': + net = ResnetDehazeGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9) + elif netG == 'unet_128': + net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout) + elif netG == 'unet_256': + net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout) + elif netG == 'unet_trans_256': + net = UnetTransGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout, r=10, eps=1e-3) + elif netG == 'haze_refine_2': + net = HazeRefiner(input_nc, output_nc, 2) + elif netG == 'haze_refine_10': + net = HazeRefiner(input_nc, output_nc, 10) + else: + raise NotImplementedError('Generator model name [%s] is not recognized' % netG) + return init_net(net, init_type, init_gain, gpu_ids) + + +def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]): + """Create a discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the first conv layer + netD (str) -- the architecture's name: basic | n_layers | pixel + n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers' + norm (str) -- the type of normalization layers used in the network. + init_type (str) -- the name of the initialization method. + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Returns a discriminator + + Our current implementation provides three types of discriminators: + [basic]: 'PatchGAN' classifier described in the original pix2pix paper. + It can classify whether 70×70 overlapping patches are real or fake. + Such a patch-level discriminator architecture has fewer parameters + than a full-image discriminator and can work on arbitrarily-sized images + in a fully convolutional fashion. + + [n_layers]: With this mode, you cna specify the number of conv layers in the discriminator + with the parameter (default=3 as used in [basic] (PatchGAN).) + + [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not. + It encourages greater color diversity but has no effect on spatial statistics. + + The discriminator has been initialized by . It uses Leakly RELU for non-linearity. + """ + net = None + norm_layer = get_norm_layer(norm_type=norm) + + if netD == 'basic': # default PatchGAN classifier + net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer) + elif netD == 'n_layers': # more options + net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer) + elif netD == 'pixel': # classify if each pixel is real or fake + net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer) + else: + raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD) + return init_net(net, init_type, init_gain, gpu_ids) + + +############################################################################## +# Classes +############################################################################## +class TVLoss(nn.Module): + ''' + Define Total Variance Loss for images + which is used for smoothness regularization + ''' + + def __init__(self): + super(TVLoss, self).__init__() + + def __call__(self, input): + # Tensor with shape (n_Batch, C, H, W) + origin = input[:, :, :-1, :-1] + right = input[:, :, :-1, 1:] + down = input[:, :, 1:, :-1] + + tv = torch.mean(torch.abs(origin-right)) + torch.mean(torch.abs(origin-down)) + return tv * 0.5 + + +class ARefiner(nn.Module): + """docstring for ARefiner""" + def __init__(self, ngf, n_downsampling): + super(ARefiner, self).__init__() + self.n_downsampling = n_downsampling + self.ngf = ngf + + self.global_pooling = torch.nn.AdaptiveAvgPool2d(1) + mult = 2 ** self.n_downsampling + self.fully_connected = torch.nn.Linear(self.ngf*mult, 3) + + def forward(self, x): + # real_I = x['real_I'] + # dcp_J = x['dcp_J'] + tmp = self.global_pooling(x) + tmp = tmp.view(tmp.shape[0],-1) + refined_A = F.relu(self.fully_connected(tmp)) + + return refined_A + + +class HazeRefiner(nn.Module): + """docstring for HazeRefiner""" + def __init__(self, input_nc, output_nc, block_num): + super(HazeRefiner, self).__init__() + self.block_num = max(2, block_num) + + self.block_0 = RefinerBlock(input_nc, output_nc) + for id_b in range(1, self.block_num): + setattr(self, 'block_%d'%id_b, RefinerBlock(input_nc+output_nc*id_b, output_nc)) + + def forward(self, x): + # real_I = x['real_I'] + # dcp_J = x['dcp_J']] + last_J = x + cur_refine_J = self.block_0(last_J) + for id_b in range(1, self.block_num): + cur_block = getattr(self, 'block_%d'%id_b) + last_J = torch.cat((last_J, cur_refine_J), 1) + cur_refine_J = cur_block(last_J) + + return cur_refine_J + + +class RefinerBlock(nn.Module): + """docstring for RefinerBlock""" + def __init__(self, input_nc, output_nc): + super(RefinerBlock, self).__init__() + + self.relu=nn.LeakyReLU(0.2, inplace=True) + + self.tanh=nn.Tanh() + + self.refine1= nn.Conv2d(input_nc, 20, kernel_size=3,stride=1,padding=1) + self.refine2= nn.Conv2d(20, 20, kernel_size=3,stride=1,padding=1) + + self.conv1010 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm + self.conv1020 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm + self.conv1030 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm + self.conv1040 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm + + self.refine3= nn.Conv2d(20+4, output_nc, kernel_size=3,stride=1,padding=1) + + # self.upsample = F.upsample_nearest + # self.batch1 = nn.InstanceNorm2d(100, affine=True) + + def forward(self, x): + output = self.relu((self.refine1(x))) + output = self.relu((self.refine2(output))) + shape_out = output.data.size() + # print(shape_out) + shape_out = shape_out[2:4] + + x101 = F.avg_pool2d(output, 32) + + x102 = F.avg_pool2d(output, 16) + + x103 = F.avg_pool2d(output, 8) + + x104 = F.avg_pool2d(output, 4) + x1010 = F.interpolate(self.relu(self.conv1010(x101)),size=shape_out, mode='nearest') + x1020 = F.interpolate(self.relu(self.conv1020(x102)),size=shape_out, mode='nearest') + x1030 = F.interpolate(self.relu(self.conv1030(x103)),size=shape_out, mode='nearest') + x1040 = F.interpolate(self.relu(self.conv1040(x104)),size=shape_out, mode='nearest') + + output = torch.cat((x1010, x1020, x1030, x1040, output), 1) + output= self.tanh(self.refine3(output)) + + return output + + +class GANLoss(nn.Module): + """Define different GAN objectives. + + The GANLoss class abstracts away the need to create the target label tensor + that has the same size as the input. + """ + + def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0): + """ Initialize the GANLoss class. + + Parameters: + gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp. + target_real_label (bool) - - label for a real image + target_fake_label (bool) - - label of a fake image + + Note: Do not use sigmoid as the last layer of Discriminator. + LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss. + """ + super(GANLoss, self).__init__() + self.register_buffer('real_label', torch.tensor(target_real_label)) + self.register_buffer('fake_label', torch.tensor(target_fake_label)) + self.gan_mode = gan_mode + if gan_mode == 'lsgan': + self.loss = nn.MSELoss() + elif gan_mode == 'vanilla': + self.loss = nn.BCEWithLogitsLoss() + elif gan_mode in ['wgangp']: + self.loss = None + else: + raise NotImplementedError('gan mode %s not implemented' % gan_mode) + + def get_target_tensor(self, prediction, target_is_real): + """Create label tensors with the same size as the input. + + Parameters: + prediction (tensor) - - tpyically the prediction from a discriminator + target_is_real (bool) - - if the ground truth label is for real images or fake images + + Returns: + A label tensor filled with ground truth label, and with the size of the input + """ + + if target_is_real: + target_tensor = self.real_label + else: + target_tensor = self.fake_label + return target_tensor.expand_as(prediction) + + def __call__(self, prediction, target_is_real): + """Calculate loss given Discriminator's output and grount truth labels. + + Parameters: + prediction (tensor) - - tpyically the prediction output from a discriminator + target_is_real (bool) - - if the ground truth label is for real images or fake images + + Returns: + the calculated loss. + """ + if self.gan_mode in ['lsgan', 'vanilla']: + target_tensor = self.get_target_tensor(prediction, target_is_real) + loss = self.loss(prediction, target_tensor) + elif self.gan_mode == 'wgangp': + if target_is_real: + loss = -prediction.mean() + else: + loss = prediction.mean() + return loss + + +def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0): + """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028 + + Arguments: + netD (network) -- discriminator network + real_data (tensor array) -- real images + fake_data (tensor array) -- generated images from the generator + device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') + type (str) -- if we mix real and fake data or not [real | fake | mixed]. + constant (float) -- the constant used in formula ( | |gradient||_2 - constant)^2 + lambda_gp (float) -- weight for this loss + + Returns the gradient penalty loss + """ + if lambda_gp > 0.0: + if type == 'real': # either use real images, fake images, or a linear interpolation of two. + interpolatesv = real_data + elif type == 'fake': + interpolatesv = fake_data + elif type == 'mixed': + alpha = torch.rand(real_data.shape[0], 1, device=device) + alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape) + interpolatesv = alpha * real_data + ((1 - alpha) * fake_data) + else: + raise NotImplementedError('{} not implemented'.format(type)) + interpolatesv.requires_grad_(True) + disc_interpolates = netD(interpolatesv) + gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv, + grad_outputs=torch.ones(disc_interpolates.size()).to(device), + create_graph=True, retain_graph=True, only_inputs=True) + gradients = gradients[0].view(real_data.size(0), -1) # flat the data + gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp # added eps + return gradient_penalty, gradients + else: + return 0.0, None + + +class ResnetGWithIntermediate(nn.Module): + """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations. + + We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style) + """ + + def __init__(self, input_nc, output_nc, ngf=6, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect', filtering='guided', r=10, eps=1e-3): + """Construct a Resnet-based generator + + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + ngf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + use_dropout (bool) -- if use dropout layers + n_blocks (int) -- the number of ResNet blocks + padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero + """ + assert(n_blocks >= 0) + super(ResnetGWithIntermediate, self).__init__() + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + self.filtering=filtering + + model = [nn.ReflectionPad2d(3), + nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), + norm_layer(ngf), + nn.ReLU(True)] + + n_downsampling = 2 + for i in range(n_downsampling): # add downsampling layers + mult = 2 ** i + model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias), + norm_layer(ngf * mult * 2), + nn.ReLU(True)] + + mult = 2 ** n_downsampling + for i in range(n_blocks): # add ResNet blocks + + model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)] + + model_up_part = [] + for i in range(n_downsampling): # add upsampling layers + mult = 2 ** (n_downsampling - i) + model_up_part += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), + kernel_size=3, stride=2, + padding=1, output_padding=1, + bias=use_bias), + norm_layer(int(ngf * mult / 2)), + nn.ReLU(True)] + model_up_part += [nn.ReflectionPad2d(3)] + model_up_part += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)] + model_up_part += [nn.Tanh()] + + self.downsampling = nn.Sequential(*model) + self.upsampling = nn.Sequential(*model_up_part) + + if self.filtering is not None: + if self.filtering == 'max': + self.last_layer = nn.MaxPool2d(kernel_size=7, stride=1, padding=3) + elif self.filtering == 'guided': + self.last_layer = GuidedFilter(r=r, eps=eps) + + def forward(self, x): + """Standard forward""" + down_out = self.downsampling(x) + up_out = self.upsampling(down_out) + # rescale to [0,1] + up_out = (up_out + 1)/2 + + # rgb2gray + guidance = 0.2989 * x[:,0,:,:] + 0.5870 * x[:,1,:,:] + 0.1140 * x[:,2,:,:] + # rescale to [0,1] + guidance = (guidance + 1) / 2 + guidance = torch.unsqueeze(guidance, dim=1) + + if up_out.shape[2:4] != guidance.shape[2:4]: + up_out = F.interpolate(up_out,size=guidance.shape[2:4], mode='nearest') + + # up_out = self.last_layer(guidance, up_out) + return self.last_layer(guidance, up_out), up_out + + +# Guided image filtering for grayscale images +class GuidedFilter(nn.Module): + def __init__(self, r=40, eps=1e-3, gpu_ids=None): # only work for gpu case at this moment + super(GuidedFilter, self).__init__() + self.r = r + self.eps = eps + # 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.boxfilter = nn.AvgPool2d(kernel_size=2*self.r+1, stride=1,padding=self.r) + + def forward(self, I, p): + """ + I -- guidance image, should be [0, 1] + p -- filtering input image, should be [0, 1] + """ + + # N = self.boxfilter(self.tensor(p.size()).fill_(1)) + N = self.boxfilter( torch.ones(p.size()) ) + + if I.is_cuda: + N = N.cuda() + + # print(N.shape) + # print(I.shape) + # print('-----------') + + mean_I = self.boxfilter(I) / N + mean_p = self.boxfilter(p) / N + mean_Ip = self.boxfilter(I*p) / N + cov_Ip = mean_Ip - mean_I * mean_p + + mean_II = self.boxfilter(I*I) / N + var_I = mean_II - mean_I * mean_I + + a = cov_Ip / (var_I + self.eps) + b = mean_p - a * mean_I + mean_a = self.boxfilter(a) / N + mean_b = self.boxfilter(b) / N + + return mean_a * I + mean_b + + +class ResnetDehazeGenerator(nn.Module): + """docstring for ResnetGenerator""" + def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'): + super(ResnetDehazeGenerator, self).__init__() + self.resnetG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, use_dropout, n_blocks, padding_type) + self.refiner = HazeRefiner(input_nc, output_nc, block_num=2) + + def forward(self, x): + res_out = self.resnetG(x) + ref_out = self.refiner(res_out) + + return ref_out + + + +class ResnetGenerator(nn.Module): + """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations. + + We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style) + """ + + def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'): + """Construct a Resnet-based generator + + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + ngf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + use_dropout (bool) -- if use dropout layers + n_blocks (int) -- the number of ResNet blocks + padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero + """ + assert(n_blocks >= 0) + super(ResnetGenerator, self).__init__() + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + model = [nn.ReflectionPad2d(3), + nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), + norm_layer(ngf), + nn.ReLU(True)] + + n_downsampling = 2 + for i in range(n_downsampling): # add downsampling layers + mult = 2 ** i + model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias), + norm_layer(ngf * mult * 2), + nn.ReLU(True)] + + mult = 2 ** n_downsampling + for i in range(n_blocks): # add ResNet blocks + + model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)] + + for i in range(n_downsampling): # add upsampling layers + mult = 2 ** (n_downsampling - i) + model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), + kernel_size=3, stride=2, + padding=1, output_padding=1, + bias=use_bias), + norm_layer(int(ngf * mult / 2)), + nn.ReLU(True)] + model += [nn.ReflectionPad2d(3)] + model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)] + model += [nn.Tanh()] + + self.model = nn.Sequential(*model) + + def forward(self, input): + """Standard forward""" + return self.model(input) + + +class ResnetBlock(nn.Module): + """Define a Resnet block""" + + def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias): + """Initialize the Resnet block + + A resnet block is a conv block with skip connections + We construct a conv block with build_conv_block function, + and implement skip connections in function. + Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf + """ + super(ResnetBlock, self).__init__() + self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias) + + def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias): + """Construct a convolutional block. + + Parameters: + dim (int) -- the number of channels in the conv layer. + padding_type (str) -- the name of padding layer: reflect | replicate | zero + norm_layer -- normalization layer + use_dropout (bool) -- if use dropout layers. + use_bias (bool) -- if the conv layer uses bias or not + + Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU)) + """ + conv_block = [] + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + + conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)] + if use_dropout: + conv_block += [nn.Dropout(0.5)] + + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)] + + return nn.Sequential(*conv_block) + + def forward(self, x): + """Forward function (with skip connections)""" + out = x + self.conv_block(x) # add skip connections + return out + + +class UnetGenerator(nn.Module): + """Create a Unet-based generator""" + + def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False): + """Construct a Unet generator + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7, + image of size 128x128 will become of size 1x1 # at the bottleneck + ngf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + + We construct the U-Net from the innermost layer to the outermost layer. + It is a recursive process. + """ + super(UnetGenerator, self).__init__() + # construct unet structure + unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True) # add the innermost layer + for i in range(num_downs - 5): # add intermediate layers with ngf * 8 filters + unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout) + # gradually reduce the number of filters from ngf * 8 to ngf + unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer) + unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer) + unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer) + self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer) # add the outermost layer + + def forward(self, input): + """Standard forward""" + return self.model(input) + + +class DCPDehazeGenerator(nn.Module): + """Create a DCP Dehaze generator""" + def __init__(self, win_size=5, r=15, eps=1e-3): + super(DCPDehazeGenerator, self).__init__() + + self.guided_filter = GuidedFilter(r=r, eps=eps) + self.neighborhood_size = win_size + self.omega = 0.95 + + def get_dark_channel(self, img, neighborhood_size): + shape = img.shape + if len(shape) == 4: + img_min,_ = torch.min(img, dim=1) + + padSize = int(np.floor(neighborhood_size/2)) + if neighborhood_size % 2 == 0: + pads = [padSize, padSize-1 ,padSize ,padSize-1] + else: + pads = [padSize, padSize ,padSize ,padSize] + + img_min = F.pad(img_min, pads, mode='constant', value=1) + dark_img = -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]') + + dark_img = torch.unsqueeze(dark_img, dim=1) + return dark_img + + def atmospheric_light(self, img, dark_img): + num,chl,height,width = img.shape + topNum = int(0.01*height*width) + + A = torch.Tensor(num,chl,1,1) + if img.is_cuda: + A = A.cuda() + + for num_id in range(num): + curImg = img[num_id,...] + curDarkImg = dark_img[num_id,0,...] + + _, indices = curDarkImg.reshape([height*width]).sort(descending=True) + #curMask = indices < topNum + + for chl_id in range(chl): + imgSlice = curImg[chl_id,...].reshape([height*width]) + A[num_id,chl_id,0,0] = torch.mean(imgSlice[indices[0:topNum]]) + + return A + + + def forward(self, x): + if x.shape[1] > 1: + # rgb2gray + guidance = 0.2989 * x[:,0,:,:] + 0.5870 * x[:,1,:,:] + 0.1140 * x[:,2,:,:] + else: + guidance = x + # rescale to [0,1] + guidance = (guidance + 1)/2 + guidance = torch.unsqueeze(guidance, dim=1) + imgPatch = (x + 1)/2 + + num,chl,height,width = imgPatch.shape + + # dark_img and A with the range of [0,1] + dark_img = self.get_dark_channel(imgPatch, self.neighborhood_size) + A = self.atmospheric_light(imgPatch, dark_img) + + map_A = A.repeat(1,1,height,width) + # make sure channel of trans_raw == 1 + trans_raw = 1 - self.omega*self.get_dark_channel(imgPatch/map_A, self.neighborhood_size) + + # get initial results + T_DCP = self.guided_filter(guidance, trans_raw) + J_DCP = (imgPatch - map_A)/T_DCP.repeat(1,3,1,1) + map_A + + # import cv2 + # temp = cv2.cvtColor(J_DCP[0].numpy().transpose([1,2,0]), cv2.COLOR_BGR2RGB) + # cv2.imshow('J_DCP',temp) + # cv2.imshow('T_DCP',T_DCP[0].numpy().transpose([1,2,0])) + # cv2.waitKey(0) + # exit() + + return J_DCP, T_DCP, torch.squeeze(A) + + +class UnetTransGenerator(nn.Module): + """Create a Unet-based generator""" + + def __init__(self, input_nc, output_nc, num_downs, ngf=6, norm_layer=nn.BatchNorm2d, use_dropout=False, r=10, eps=1e-3): + """Construct a Unet generator + Parameters: + input_nc (int) -- the number of channels in input images + output_nc (int) -- the number of channels in output images + num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7, + image of size 128x128 will become of size 1x1 # at the bottleneck + ngf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + + We construct the U-Net from the innermost layer to the outermost layer. + It is a recursive process. + """ + super(UnetTransGenerator, self).__init__() + # construct unet structure + unet_block = UnetAlignedSkipBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True) # add the innermost layer + for i in range(num_downs - 5): # add intermediate layers with ngf * 8 filters + unet_block = UnetAlignedSkipBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout) + # gradually reduce the number of filters from ngf * 8 to ngf + unet_block = UnetAlignedSkipBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer) + unet_block = UnetAlignedSkipBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer) + unet_block = UnetAlignedSkipBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer) + self.model = UnetAlignedSkipBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer) # add the outermost layer + + self.guided_filter = GuidedFilter(r=r, eps=eps) + + def forward(self, x): + if x.shape[1] > 1: + # rgb2gray + guidance = 0.2989 * x[:,0,:,:] + 0.5870 * x[:,1,:,:] + 0.1140 * x[:,2,:,:] + else: + guidance = x + # rescale to [0,1] + guidance = (guidance + 1) / 2 + guidance = torch.unsqueeze(guidance, dim=1) + + trans_raw = (self.model(x) + 1) / 2 # transmission ranges [0,1] + if trans_raw.shape[2:4] != guidance.shape[2:4]: + trans_raw = F.interpolate(trans_raw,size=guidance.shape[2:4], mode='nearest') + + return self.guided_filter(guidance, trans_raw), trans_raw + +# for trans refination +class UnetAlignedSkipBlock(nn.Module): + """Defines the Unet submodule with skip connection. + X -------------------identity---------------------- + |-- downsampling -- |submodule| -- upsampling --| + """ + + def __init__(self, outer_nc, inner_nc, input_nc=None, + submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False): + """Construct a Unet submodule with skip connections. + + Parameters: + outer_nc (int) -- the number of filters in the outer conv layer + inner_nc (int) -- the number of filters in the inner conv layer + input_nc (int) -- the number of channels in input images/features + submodule (UnetAlignedSkipBlock) -- previously defined submodules + outermost (bool) -- if this module is the outermost module + innermost (bool) -- if this module is the innermost module + norm_layer -- normalization layer + user_dropout (bool) -- if use dropout layers. + """ + super(UnetAlignedSkipBlock, self).__init__() + self.outermost = outermost + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + if input_nc is None: + input_nc = outer_nc + downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, + stride=2, padding=1, bias=use_bias) + downrelu = nn.LeakyReLU(0.2, True) + downnorm = norm_layer(inner_nc) + uprelu = nn.ReLU(True) + upnorm = norm_layer(outer_nc) + + if outermost: + upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, + kernel_size=4, stride=2, + padding=1) + down = [downconv] + up = [uprelu, upconv, nn.Tanh()] + model = down + [submodule] + up + elif innermost: + upconv = nn.ConvTranspose2d(inner_nc, outer_nc, + kernel_size=4, stride=2, + padding=1, bias=use_bias) + down = [downrelu, downconv] + up = [uprelu, upconv, upnorm] + model = down + up + else: + upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, + kernel_size=4, stride=2, + padding=1, bias=use_bias) + down = [downrelu, downconv, downnorm] + up = [uprelu, upconv, upnorm] + + if use_dropout: + model = down + [submodule] + up + [nn.Dropout(0.5)] + else: + model = down + [submodule] + up + + self.model = nn.Sequential(*model) + + def forward(self, x): + if self.outermost: + return self.model(x) + else: # add skip connections + y = self.model(x) + # print(x.shape, y.shape) + if x.shape != y.shape: + y = F.interpolate(y, size=x.shape[2:4], mode='nearest') + return torch.cat([x, y], 1) + +class UnetSkipConnectionBlock(nn.Module): + """Defines the Unet submodule with skip connection. + X -------------------identity---------------------- + |-- downsampling -- |submodule| -- upsampling --| + """ + + def __init__(self, outer_nc, inner_nc, input_nc=None, + submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False): + """Construct a Unet submodule with skip connections. + + Parameters: + outer_nc (int) -- the number of filters in the outer conv layer + inner_nc (int) -- the number of filters in the inner conv layer + input_nc (int) -- the number of channels in input images/features + submodule (UnetSkipConnectionBlock) -- previously defined submodules + outermost (bool) -- if this module is the outermost module + innermost (bool) -- if this module is the innermost module + norm_layer -- normalization layer + user_dropout (bool) -- if use dropout layers. + """ + super(UnetSkipConnectionBlock, self).__init__() + self.outermost = outermost + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + if input_nc is None: + input_nc = outer_nc + downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, + stride=2, padding=1, bias=use_bias) + downrelu = nn.LeakyReLU(0.2, True) + downnorm = norm_layer(inner_nc) + uprelu = nn.ReLU(True) + upnorm = norm_layer(outer_nc) + + if outermost: + upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, + kernel_size=4, stride=2, + padding=1) + down = [downconv] + up = [uprelu, upconv, nn.Tanh()] + model = down + [submodule] + up + elif innermost: + upconv = nn.ConvTranspose2d(inner_nc, outer_nc, + kernel_size=4, stride=2, + padding=1, bias=use_bias) + down = [downrelu, downconv] + up = [uprelu, upconv, upnorm] + model = down + up + else: + upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, + kernel_size=4, stride=2, + padding=1, bias=use_bias) + down = [downrelu, downconv, downnorm] + up = [uprelu, upconv, upnorm] + + if use_dropout: + model = down + [submodule] + up + [nn.Dropout(0.5)] + else: + model = down + [submodule] + up + + self.model = nn.Sequential(*model) + + def forward(self, x): + if self.outermost: + return self.model(x) + else: # add skip connections + return torch.cat([x, self.model(x)], 1) + + +class NLayerDiscriminator(nn.Module): + """Defines a PatchGAN discriminator""" + + def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d): + """Construct a PatchGAN discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + n_layers (int) -- the number of conv layers in the discriminator + norm_layer -- normalization layer + """ + super(NLayerDiscriminator, self).__init__() + if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + kw = 4 + padw = 1 + sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)] + nf_mult = 1 + nf_mult_prev = 1 + for n in range(1, n_layers): # gradually increase the number of filters + nf_mult_prev = nf_mult + nf_mult = min(2 ** n, 8) + sequence += [ + nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + nf_mult_prev = nf_mult + nf_mult = min(2 ** n_layers, 8) + sequence += [ + nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map + self.model = nn.Sequential(*sequence) + + def forward(self, input): + """Standard forward.""" + return self.model(input) + + +class PixelDiscriminator(nn.Module): + """Defines a 1x1 PatchGAN discriminator (pixelGAN)""" + + def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d): + """Construct a 1x1 PatchGAN discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + """ + super(PixelDiscriminator, self).__init__() + if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + self.net = [ + nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0), + nn.LeakyReLU(0.2, True), + nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias), + norm_layer(ndf * 2), + nn.LeakyReLU(0.2, True), + nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)] + + self.net = nn.Sequential(*self.net) + + def forward(self, input): + """Standard forward.""" + return self.net(input) + + +# Defines the Multiscale-PatchGAN discriminator with the specified arguments. +class MultiDiscriminator(nn.Module): + def __init__(self, input_nc, ndf=64, n_layers=5, norm_layer=nn.BatchNorm2d, use_sigmoid=False, gpu_ids=[]): + super(MultiDiscriminator, self).__init__() + self.gpu_ids = gpu_ids + if type(norm_layer) == functools.partial: + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + # cannot deal with use_sigmoid=True case at thie moment + assert(use_sigmoid == False) + + kw = 4 + padw = int(np.ceil((kw-1)/2)) + scale1 = [ + nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, True) + ] + + nf_mult = 1 + nf_mult_prev = 1 + for n in range(1, 3): + nf_mult_prev = nf_mult + nf_mult = min(2**n, 8) + scale1 += [ + nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, + kernel_size=kw, stride=2, padding=padw, bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + self.scale1 = nn.Sequential(*scale1) + scale1_output = [] + scale1_output += [ + nn.Conv2d(ndf * nf_mult, ndf * nf_mult, + kernel_size=kw, stride=1, padding=padw, bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + scale1_output += [nn.Conv2d(ndf*nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # compress to 1 channel + self.scale1_output = nn.Sequential(*scale1_output) + + scale2 = [] + nf_mult = nf_mult + for n in range(3, n_layers): + nf_mult_prev = nf_mult + nf_mult = min(2**n, 8) + scale2 += [ + nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, + kernel_size=kw, stride=2, padding=padw, bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + nf_mult_prev = nf_mult + nf_mult = min(2**n_layers, 8) + scale2 += [ + nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, + kernel_size=kw, stride=1, padding=padw, bias=use_bias), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True) + ] + + scale2 += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] + + if use_sigmoid: + scale2 += [nn.Sigmoid()] + + self.scale2 = nn.Sequential(*scale2) + + def forward(self, input): + if len(self.gpu_ids) and isinstance(input.data, torch.cuda.FloatTensor): + scale1 = nn.parallel.data_parallel(self.scale1, input, self.gpu_ids) + output1 = nn.parallel.data_parallel(self.scale1_output, scale1, self.gpu_ids) + output2 = nn.parallel.data_parallel(self.scale2, scale1, self.gpu_ids) + else: + scale1 = self.scale1(input) + output1 = self.scale1_output(scale1) + output2 = self.scale2(scale1) + + return output1, output2 + +class VGGLoss(nn.Module): + def __init__(self): + super(VGGLoss, self).__init__() + self.vgg = Vgg19().cuda() + self.criterion = nn.L1Loss() + self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0] + # self.weights = [0, 0, 1.0/8, 1.0/4, 1.0] + + def forward(self, x, y): + x_vgg, y_vgg = self.vgg(x), self.vgg(y) + loss = 0 + for i in range(len(x_vgg)): + loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach()) + return loss + + +class Vgg19(torch.nn.Module): + def __init__(self, requires_grad=False): + super(Vgg19, self).__init__() + from torchvision import models + vgg_pretrained_features = models.vgg19(pretrained=True).features + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + self.slice4 = torch.nn.Sequential() + self.slice5 = torch.nn.Sequential() + for x in range(2): + self.slice1.add_module(str(x), vgg_pretrained_features[x]) + for x in range(2, 7): + self.slice2.add_module(str(x), vgg_pretrained_features[x]) + for x in range(7, 12): + self.slice3.add_module(str(x), vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice4.add_module(str(x), vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice5.add_module(str(x), vgg_pretrained_features[x]) + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + h_relu1 = self.slice1(X) + h_relu2 = self.slice2(h_relu1) + h_relu3 = self.slice3(h_relu2) + h_relu4 = self.slice4(h_relu3) + h_relu5 = self.slice5(h_relu4) + out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] + return out diff --git a/RefineDNet/models/pix2pix_model.py b/RefineDNet/models/pix2pix_model.py new file mode 100644 index 0000000..939eb88 --- /dev/null +++ b/RefineDNet/models/pix2pix_model.py @@ -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 + 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 + 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 and + 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 . + 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 and .""" + 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 diff --git a/RefineDNet/models/refined_DCP_model.py b/RefineDNet/models/refined_DCP_model.py new file mode 100644 index 0000000..bc7456f --- /dev/null +++ b/RefineDNet/models/refined_DCP_model.py @@ -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 + 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 + 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 and . + 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 . + 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 and .""" + 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 function in no_grad() so we don't save intermediate steps for backprop + It also calls 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 diff --git a/RefineDNet/models/template_model.py b/RefineDNet/models/template_model.py new file mode 100644 index 0000000..68cdaf6 --- /dev/null +++ b/RefineDNet/models/template_model.py @@ -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 _dataset.py +The class name should be 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(data_A) - data_B||_1 +You need to implement the following functions: + :　Add model-specific options and rewrite default values for existing options. + <__init__>: Initialize this model class. + : Unpack input data and perform data pre-processing. + : Run forward pass. This will be called by both and . + : 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 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 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 and .""" + 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 + # 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 diff --git a/RefineDNet/models/test_model.py b/RefineDNet/models/test_model.py new file mode 100644 index 0000000..fe15f40 --- /dev/null +++ b/RefineDNet/models/test_model.py @@ -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 + self.loss_names = [] + # specify the images you want to save/display. The training/test scripts will call + self.visual_names = ['real', 'fake'] + # specify the models you want to save to the disk. The training/test scripts will call and + 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 + 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 diff --git a/RefineDNet/options/__init__.py b/RefineDNet/options/__init__.py new file mode 100644 index 0000000..e7eedeb --- /dev/null +++ b/RefineDNet/options/__init__.py @@ -0,0 +1 @@ +"""This package options includes option modules: training options, test options, and basic options (used in both training and test).""" diff --git a/RefineDNet/options/base_options.py b/RefineDNet/options/base_options.py new file mode 100644 index 0000000..6c212b3 --- /dev/null +++ b/RefineDNet/options/base_options.py @@ -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 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 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 diff --git a/RefineDNet/options/test_options.py b/RefineDNet/options/test_options.py new file mode 100644 index 0000000..660247b --- /dev/null +++ b/RefineDNet/options/test_options.py @@ -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 diff --git a/RefineDNet/options/train_options.py b/RefineDNet/options/train_options.py new file mode 100644 index 0000000..8b8ebfb --- /dev/null +++ b/RefineDNet/options/train_options.py @@ -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 , +, ...') + 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 diff --git a/RefineDNet/quick_test.py b/RefineDNet/quick_test.py new file mode 100644 index 0000000..7e7bd3c --- /dev/null +++ b/RefineDNet/quick_test.py @@ -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))) diff --git a/RefineDNet/test_BeDDE.py b/RefineDNet/test_BeDDE.py new file mode 100644 index 0000000..901975c --- /dev/null +++ b/RefineDNet/test_BeDDE.py @@ -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))) diff --git a/RefineDNet/train.py b/RefineDNet/train.py new file mode 100644 index 0000000..ffe1026 --- /dev/null +++ b/RefineDNet/train.py @@ -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_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 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 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. diff --git a/RefineDNet/util/__init__.py b/RefineDNet/util/__init__.py new file mode 100644 index 0000000..ae36f63 --- /dev/null +++ b/RefineDNet/util/__init__.py @@ -0,0 +1 @@ +"""This package includes a miscellaneous collection of useful helper functions.""" diff --git a/RefineDNet/util/get_data.py b/RefineDNet/util/get_data.py new file mode 100644 index 0000000..97edc3c --- /dev/null +++ b/RefineDNet/util/get_data.py @@ -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) diff --git a/RefineDNet/util/html.py b/RefineDNet/util/html.py new file mode 100644 index 0000000..cc3262a --- /dev/null +++ b/RefineDNet/util/html.py @@ -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 a text header to the HTML file), + (add a row of images to the HTML file), and (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 /index.html; images will be saved at 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() diff --git a/RefineDNet/util/image_pool.py b/RefineDNet/util/image_pool.py new file mode 100644 index 0000000..6d086f8 --- /dev/null +++ b/RefineDNet/util/image_pool.py @@ -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 diff --git a/RefineDNet/util/util.py b/RefineDNet/util/util.py new file mode 100644 index 0000000..296ab8a --- /dev/null +++ b/RefineDNet/util/util.py @@ -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) diff --git a/RefineDNet/util/visualizer.py b/RefineDNet/util/visualizer.py new file mode 100644 index 0000000..b45169c --- /dev/null +++ b/RefineDNet/util/visualizer.py @@ -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 and + 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 /web/; images will be saved under /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 = """""" % (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 += '%s' % label + images.append(image_numpy.transpose([2, 0, 1])) + idx += 1 + if idx % ncols == 0: + label_html += '%s' % 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 += '' + idx += 1 + if label_html_row != '': + label_html += '%s' % label_html_row + try: + self.vis.images(images, nrow=ncols, win=self.display_id + 1, + padding=2, opts=dict(title=title + ' images')) + label_html = '%s
' % 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 diff --git a/config/baidu_api.env.example b/config/baidu_api.env.example new file mode 100644 index 0000000..ab9838d --- /dev/null +++ b/config/baidu_api.env.example @@ -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 diff --git a/run_dehaze_web.sh b/run_dehaze_web.sh new file mode 100755 index 0000000..cd19634 --- /dev/null +++ b/run_dehaze_web.sh @@ -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 diff --git a/scripts/clean_generated.sh b/scripts/clean_generated.sh new file mode 100755 index 0000000..eccd69c --- /dev/null +++ b/scripts/clean_generated.sh @@ -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." diff --git a/scripts/verify_all.py b/scripts/verify_all.py new file mode 100755 index 0000000..45482df --- /dev/null +++ b/scripts/verify_all.py @@ -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()) diff --git a/web_dehaze/README.md b/web_dehaze/README.md new file mode 100644 index 0000000..5036236 --- /dev/null +++ b/web_dehaze/README.md @@ -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` diff --git a/web_dehaze/pipeline.py b/web_dehaze/pipeline.py new file mode 100644 index 0000000..9cdb2c6 --- /dev/null +++ b/web_dehaze/pipeline.py @@ -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 diff --git a/web_dehaze/postprocess.py b/web_dehaze/postprocess.py new file mode 100644 index 0000000..29e5c12 --- /dev/null +++ b/web_dehaze/postprocess.py @@ -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}") diff --git a/web_dehaze/server.py b/web_dehaze/server.py new file mode 100644 index 0000000..ad6bdd6 --- /dev/null +++ b/web_dehaze/server.py @@ -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() diff --git a/web_dehaze/static/app.js b/web_dehaze/static/app.js new file mode 100644 index 0000000..0e0b0a2 --- /dev/null +++ b/web_dehaze/static/app.js @@ -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 = ` + ${image.name} + ${image.width || "?"}x${image.height || "?"} · ${formatBytes(image.size)} + `; + 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 = ` + + ${info.label} + ${info.available ? "ready" : `缺 ${info.module_ok ? "模型" : info.module}`} + + + `; + 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 = ` + ${info.label} + + `; + 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 ? 'ready' : '未生成'; + const body = exists + ? `
${title}
` + : '
暂无结果
'; + card.innerHTML = `

${title}

${badge}
${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; +}); diff --git a/web_dehaze/static/index.html b/web_dehaze/static/index.html new file mode 100644 index 0000000..1a3c8de --- /dev/null +++ b/web_dehaze/static/index.html @@ -0,0 +1,93 @@ + + + + + + Dehaze Console + + + +
+ + +
+
+
+

当前图片

+

未选择

+
+
Idle
+
+ +
+ +
+
+

日志

+ +
+ 

+        

+      

+    

+
+    
+  
+
diff --git a/web_dehaze/static/style.css b/web_dehaze/static/style.css
new file mode 100644
index 0000000..38ecb1b
--- /dev/null
+++ b/web_dehaze/static/style.css
@@ -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;
+  }
+}
diff --git a/※程序-后处理汇总/2_Adjust_V1_HSV中SV直方图调整_很一般.py b/※程序-后处理汇总/2_Adjust_V1_HSV中SV直方图调整_很一般.py
new file mode 100644
index 0000000..376c9f5
--- /dev/null
+++ b/※程序-后处理汇总/2_Adjust_V1_HSV中SV直方图调整_很一般.py
@@ -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")
\ No newline at end of file
diff --git a/※程序-后处理汇总/2_Adjust_V2_SV自动调整_※Best.py b/※程序-后处理汇总/2_Adjust_V2_SV自动调整_※Best.py
new file mode 100644
index 0000000..ceccfd1
--- /dev/null
+++ b/※程序-后处理汇总/2_Adjust_V2_SV自动调整_※Best.py
@@ -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("错误: 找不到输入文件夹或参考文件夹，请检查路径。")
\ No newline at end of file
diff --git a/※程序-后处理汇总/2_Adjust_V3_SV直方图调整+SV自动调整_效果一般.py b/※程序-后处理汇总/2_Adjust_V3_SV直方图调整+SV自动调整_效果一般.py
new file mode 100644
index 0000000..9c49821
--- /dev/null
+++ b/※程序-后处理汇总/2_Adjust_V3_SV直方图调整+SV自动调整_效果一般.py
@@ -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()
\ No newline at end of file
diff --git a/使用手册.md b/使用手册.md
new file mode 100644
index 0000000..e93f2a5
--- /dev/null
+++ b/使用手册.md
@@ -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 模型、权重与入口脚本。
+- `※程序-后处理汇总/`：原始后处理脚本归档。
diff --git a/待去雾图片/1.png b/待去雾图片/1.png
new file mode 100644
index 0000000..57382ef
Binary files /dev/null and b/待去雾图片/1.png differ
diff --git a/待去雾图片/2.jpg b/待去雾图片/2.jpg
new file mode 100644
index 0000000..643d4cc
Binary files /dev/null and b/待去雾图片/2.jpg differ