整合去雾网页工具

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}
+}
+```

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

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)

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