Upload 7 files

models/__init__.py

@@ -0,0 +1,2 @@
+from .networks import *
+

models/losses.py

@@ -0,0 +1,279 @@
+import torch
+import torch.nn as nn
+from pytorch_msssim import SSIM, MS_SSIM
+from torch.nn import L1Loss, MSELoss
+from torchvision.models import vgg16
+import torch.nn.functional as F
+
+
+def compute_gradient(img):
+    gradx = img[..., 1:, :] - img[..., :-1, :]
+    grady = img[..., 1:] - img[..., :-1]
+    return gradx, grady
+
+
+class GradientLoss(nn.Module):
+    def __init__(self):
+        super(GradientLoss, self).__init__()
+        self.loss = nn.L1Loss()
+
+    def forward(self, predict, target):
+        predict_gradx, predict_grady = compute_gradient(predict)
+        target_gradx, target_grady = compute_gradient(target)
+
+        return self.loss(predict_gradx, target_gradx) + self.loss(predict_grady, target_grady)
+
+
+class SSIMLoss(nn.Module):
+    def __init__(self, channels):
+        super(SSIMLoss, self).__init__()
+        self.ssim = SSIM(data_range=1., size_average=True, channel=channels)
+
+    def forward(self, output, target):
+        ssim_loss = 1 - self.ssim(output, target)
+        return ssim_loss
+
+
+class SSIML1Loss(nn.Module):
+    def __init__(self, channels):
+        super(SSIML1Loss, self).__init__()
+        self.l1_loss_func = nn.L1Loss()
+        self.ssim = SSIM(data_range=1., size_average=True, channel=channels)
+        self.alpha = 1.4
+
+    def forward(self, output, target):
+        l1_loss = self.l1_loss_func(output, target)
+        ssim_loss = 1 - self.ssim(output, target)
+        total_loss = l1_loss + self.alpha * ssim_loss
+        return total_loss
+
+
+class GradSSIML1Loss(nn.Module):
+    def __init__(self, channels):
+        super(GradSSIML1Loss, self).__init__()
+        self.l1_loss_func = nn.L1Loss()
+        self.ssim = SSIM(data_range=1., size_average=True, channel=channels)
+        self.grad_loss_func = GradientLoss()
+        self.alpha = 1.4
+
+    def forward(self, output, target):
+        l1_loss = self.l1_loss_func(output, target)
+        ssim_loss = 1 - self.ssim(output, target)
+        grad_loss = self.grad_loss_func(output, target)
+        total_loss = l1_loss + self.alpha * ssim_loss + 0.2 * grad_loss
+        return total_loss
+
+
+class SSIML2Loss(nn.Module):
+    def __init__(self, channels):
+        super(SSIML2Loss, self).__init__()
+        self.l2_loss_func = nn.MSELoss()
+        self.ssim = SSIM(data_range=1., size_average=True, channel=channels)
+        self.alpha = 1.
+
+    def forward(self, output, target):
+        l2_loss = self.l2_loss_func(output, target)
+        ssim_loss = 1 - self.ssim(output, target)
+        total_loss = l2_loss + self.alpha * ssim_loss
+        return total_loss
+
+
+class MSSSIML1Loss(nn.Module):
+    def __init__(self, channels):
+        super(MSSSIML1Loss, self).__init__()
+        self.l1_loss_func = nn.L1Loss()
+        self.ms_ssim = MS_SSIM(data_range=1., size_average=True, channel=channels)
+        self.alpha = 1.0
+
+    def forward(self, output, target):
+        ms_ssim_loss = 1 - self.ms_ssim(output, target)
+        l1_loss = self.l1_loss_func(output, target)
+        total_loss = l1_loss + self.alpha * ms_ssim_loss
+        return total_loss
+
+
+class MSSSIML2Loss(nn.Module):
+    def __init__(self, channels):
+        super(MSSSIML2Loss, self).__init__()
+        self.l2_loss_func = nn.MSELoss()
+        self.ms_ssim = MS_SSIM(data_range=1., size_average=True, channel=channels)
+        # self.alpha = 0.84
+        self.alpha = 1.2
+
+    def forward(self, output, target):
+        l2_loss = self.l2_loss_func(output, target)
+        ms_ssim_loss = 1 - self.ms_ssim(output, target)
+        total_loss = l2_loss + self.alpha * ms_ssim_loss
+        return total_loss
+
+
+class PerLoss(torch.nn.Module):
+    def __init__(self):
+        super(PerLoss, self).__init__()
+        vgg_model = vgg16(pretrained=True).features[:16]
+        vgg_model = vgg_model.to('cuda')
+        for param in vgg_model.parameters():
+            param.requires_grad = False
+
+        self.vgg_layers = vgg_model
+
+        self.layer_name_mapping = {
+            '3': "relu1_2",
+            '8': "relu2_2",
+            '15': "relu3_3"
+        }
+
+    def output_features(self, x):
+        output = {}
+        for name, module in self.vgg_layers._modules.items():
+            x = module(x)
+            if name in self.layer_name_mapping:
+                output[self.layer_name_mapping[name]] = x
+        return list(output.values())
+
+    def forward(self, data, gt):
+        loss = []
+        if data.shape[1] == 1:
+            data = data.repeat(1, 3, 1, 1)
+            gt = gt.repeat(1, 3, 1, 1)
+
+        dehaze_features = self.output_features(data)
+        gt_features = self.output_features(gt)
+        for dehaze_feature, gt_feature in zip(dehaze_features, gt_features):
+            loss.append(F.mse_loss(dehaze_feature, gt_feature))
+
+        return sum(loss) / len(loss)
+
+
+class PerL1Loss(torch.nn.Module):
+    def __init__(self):
+        super(PerL1Loss, self).__init__()
+        self.l1_loss_func = nn.L1Loss()
+        self.per_loss_func = PerLoss().to('cuda')
+
+    def forward(self, output, target):
+        l1_loss = self.l1_loss_func(output, target)
+        per_loss = self.per_loss_func(output, target)
+        # total_loss = l1_loss + 0.04 * per_loss
+        total_loss = l1_loss + 0.2 * per_loss
+        return total_loss
+
+
+class MSPerL1Loss(torch.nn.Module):
+    def __init__(self, channels):
+        super(MSPerL1Loss, self).__init__()
+        self.l1_loss_func = nn.L1Loss()
+        self.ms_ssim = MS_SSIM(data_range=1., size_average=True, channel=channels)
+        self.per_loss_func = PerLoss().to('cuda')
+
+    def forward(self, output, target):
+        ms_ssim_loss = 1 - self.ms_ssim(output, target)
+        l1_loss = self.l1_loss_func(output, target)
+        per_loss = self.per_loss_func(output, target)
+        total_loss = l1_loss + 1.2 * ms_ssim_loss + 0.04 * per_loss
+        return total_loss
+
+
+class MSPerL2Loss(torch.nn.Module):
+    def __init__(self):
+        super(MSPerL2Loss, self).__init__()
+        self.l2_loss_func = nn.MSELoss()
+        self.ms_ssim = MS_SSIM(data_range=1., size_average=True, channel=3)
+        self.per_loss_func = PerLoss().to('cuda')
+
+    def forward(self, output, target):
+        ms_ssim_loss = 1 - self.ms_ssim(output, target)
+        l2_loss = self.l2_loss_func(output, target)
+        per_loss = self.per_loss_func(output, target)
+        total_loss = l2_loss + 0.16 * ms_ssim_loss + 0.2 * per_loss
+        return total_loss
+
+
+class TVLoss(torch.nn.Module):
+    def __init__(self):
+        super(TVLoss, self).__init__()
+
+    def forward(self, data):
+        w_variance = torch.sum(torch.pow(data[:, :, :, :-1] - data[:, :, :, 1:], 2))
+        h_variance = torch.sum(torch.pow(data[:, :, :-1, :] - data[:, :, 1:, :], 2))
+
+        count_h = self._tensor_size(data[:, :, 1:, :])
+        count_w = self._tensor_size(data[:, :, :, 1:])
+
+        tv_loss = h_variance / count_h + w_variance / count_w
+        return tv_loss
+
+    def _tensor_size(self, t):
+        return t.size()[1] * t.size()[2] * t.size()[3]
+
+
+def safe_div(a, b, eps=1e-2):
+    return a / torch.clamp_min(b, eps)
+
+
+class WTVLoss(torch.nn.Module):
+    def __init__(self):
+        super(WTVLoss, self).__init__()
+        self.eps = 1e-2
+
+    def forward(self, data, aux):
+        data_dw = data[:, :, :, :-1] - data[:, :, :, 1:]
+        data_dh = data[:, :, :-1, :] - data[:, :, 1:, :]
+        aux_dw = torch.abs(aux[:, :, :, :-1] - aux[:, :, :, 1:])
+        aux_dh = torch.abs(aux[:, :, :-1, :] - aux[:, :, 1:, :])
+
+        w_variance = torch.sum(torch.pow(safe_div(data_dw, aux_dw, self.eps), 2))
+        h_variance = torch.sum(torch.pow(safe_div(data_dh, aux_dh, self.eps), 2))
+
+        count_h = self._tensor_size(data[:, :, 1:, :])
+        count_w = self._tensor_size(data[:, :, :, 1:])
+
+        tv_loss = h_variance / count_h + w_variance / count_w
+        return tv_loss
+
+    def _tensor_size(self, t):
+        return t.size()[1] * t.size()[2] * t.size()[3]
+
+
+class WTVLoss2(torch.nn.Module):
+    def __init__(self):
+        super(WTVLoss2, self).__init__()
+        self.eps = 1e-2
+        self.criterion = nn.MSELoss()
+
+    def forward(self, data, aux):
+        N, C, H, W = data.shape
+
+        data_dw = F.pad(torch.abs(data[:, :, :, :-1] - data[:, :, :, 1:]), (1, 0, 0, 0))
+        data_dh = F.pad(torch.abs(data[:, :, :-1, :] - data[:, :, 1:, :]), (0, 0, 1, 0))
+        aux_dw = F.pad(torch.abs(aux[:, :, :, :-1] - aux[:, :, :, 1:]), (1, 0, 0, 0))
+        aux_dh = F.pad(torch.abs(aux[:, :, :-1, :] - aux[:, :, 1:, :]), (0, 0, 1, 0))
+
+        data_d = data_dw + data_dh
+        aux_d = aux_dw + aux_dh
+
+        loss1 = self.criterion(data_d, aux_d)
+        # loss2 = torch.norm(data_d / (aux_d + self.eps), p=1) / (C * H * W)
+        loss2 = torch.norm(data_d / (aux_d + self.eps)) / (C * H * W)
+        return loss1 * 0.5 + loss2 * 4.0
+
+
+class MSTVPerL1Loss(torch.nn.Module):
+    def __init__(self):
+        super(MSTVPerL1Loss, self).__init__()
+        self.l1_loss_func = nn.L1Loss()
+        self.ms_ssim = MS_SSIM(data_range=1., size_average=True, channel=3)
+        self.per_loss_func = PerLoss().to('cuda')
+        self.tv_loss_func = TVLoss()
+
+    def forward(self, output, target):
+        ms_ssim_loss = 1 - self.ms_ssim(output, target)
+        l1_loss = self.l1_loss_func(output, target)
+        per_loss = self.per_loss_func(output, target)
+        tv_loss = self.tv_loss_func(output)
+        total_loss = l1_loss + 1.2 * ms_ssim_loss + 0.04 * per_loss + 1e-7 * tv_loss
+        return total_loss
+
+
+if __name__ == "__main__":
+    MSTVPerL1Loss()

models/lr_scheduler.py

@@ -0,0 +1,20 @@
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+import math
+
+
+class CosineLR(_LRScheduler):
+    def __init__(self, optimizer, init_lr, total_epochs, last_epoch=-1):
+        super(CosineLR, self).__init__(optimizer, last_epoch=-1)
+        self.optimizer = optimizer
+        self.init_lr = init_lr
+        self.total_epochs = total_epochs
+        self.last_epoch = last_epoch
+        print(f'CosineLR start from epoch(step) {last_epoch} with init_lr {init_lr} ')
+
+    def get_lr(self):
+        if self.last_epoch == 0:
+            return [group['lr'] for group in self.optimizer.param_groups]
+
+        return [0.5 * (1 + math.cos(self.last_epoch * math.pi / self.total_epochs)) * self.init_lr for group in
+                self.optimizer.param_groups]

models/metrics.py

@@ -0,0 +1,63 @@
+import math
+from math import exp
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+
+def _ssim(img1, img2, window, window_size, channel, size_average=True):
+    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+
+def SSIM(img1, img2, window_size=11, size_average=True):
+    img1 = torch.clamp(img1, min=0, max=1)
+    img2 = torch.clamp(img2, min=0, max=1)
+    (_, channel, _, _) = img1.size()
+    window = create_window(window_size, channel)
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+    return _ssim(img1, img2, window, window_size, channel, size_average)
+
+
+def PSNR(pred, gt):
+    pred = pred.clamp(0, 1).detach().cpu().numpy()
+    gt = gt.clamp(0, 1).detach().cpu().numpy()
+    imdff = pred - gt
+    rmse = math.sqrt(np.mean(imdff ** 2))
+    if rmse == 0:
+        return 100
+    return 20 * math.log10(1.0 / rmse)
+
+
+if __name__ == "__main__":
+    pass

models/networks/__init__.py

@@ -0,0 +1 @@
+from .networks import *

models/networks/modules.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PALayer(nn.Module):
+    def __init__(self, channel):
+        super(PALayer, self).__init__()
+        self.pa = nn.Sequential(
+            nn.Conv2d(channel, channel // 8, 1, padding=0, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // 8, 1, 1, padding=0, bias=True),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.pa(x)
+        return x * y
+
+
+class CALayer(nn.Module):
+    def __init__(self, channel):
+        super(CALayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.ca = nn.Sequential(
+            nn.Conv2d(channel, channel // 8, 1, padding=0, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // 8, channel, 1, padding=0, bias=True),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.ca(y)
+
+        return x * y
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels, norm=False, leaky=True):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels) if norm else nn.Identity(),
+            nn.LeakyReLU(0.2, inplace=True) if leaky else nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels) if norm else nn.Identity(),
+            nn.LeakyReLU(0.2, inplace=True) if leaky else nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels, act=True):
+        super(OutConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.Sigmoid() if act else nn.Identity()
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels, norm=True, leaky=True):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels, norm=norm, leaky=leaky)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True, norm=True, leaky=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, norm=norm, leaky=leaky)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels, norm=norm, leaky=leaky)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class AttentiveDown(nn.Module):
+    def __init__(self, in_channels, out_channels, norm=False, leaky=True):
+        super().__init__()
+        self.down = Down(in_channels, out_channels, norm=norm, leaky=leaky)
+        self.attention = nn.Sequential(
+            CALayer(out_channels),
+            PALayer(out_channels)
+        )
+
+    def forward(self, x):
+        return self.attention(self.down(x))
+
+
+class AttentiveUp(nn.Module):
+    def __init__(self, in_channels, out_channels, bilinear=True, norm=False, leaky=True):
+        super().__init__()
+        self.up = Up(in_channels, out_channels, bilinear, norm=norm, leaky=leaky)
+        self.attention = nn.Sequential(
+            CALayer(out_channels),
+            PALayer(out_channels)
+        )
+
+    def forward(self, x1, x2):
+        return self.attention(self.up(x1, x2))
+
+
+class AttentiveDoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels, norm=False, leaky=False):
+        super().__init__()
+        self.conv = DoubleConv(in_channels, out_channels, norm=norm, leaky=leaky)
+        self.attention = nn.Sequential(
+            CALayer(out_channels),
+            PALayer(out_channels)
+        )
+
+    def forward(self, x):
+        return self.attention(self.conv(x))

models/networks/networks.py

@@ -0,0 +1,65 @@
+from models.networks.modules import *
+
+
+class BaseNet(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1, norm=True):
+        super(BaseNet, self).__init__()
+        self.n_channels = in_channels
+        self.n_classes = out_channels
+
+        self.inc = DoubleConv(in_channels, 32, norm=norm)
+        self.down1 = Down(32, 64, norm=norm)
+        self.down2 = Down(64, 128, norm=norm)
+        self.down3 = Down(128, 128, norm=norm)
+
+        self.up1 = Up(256, 64, bilinear=True, norm=norm)
+        self.up2 = Up(128, 32, bilinear=True, norm=norm)
+        self.up3 = Up(64, 32, bilinear=True, norm=norm)
+        self.outc = OutConv(32, out_channels)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x = self.up1(x4, x3)
+        x = self.up2(x, x2)
+        x = self.up3(x, x1)
+        logits = self.outc(x)
+        return logits
+
+
+class IAN(BaseNet):
+    def __init__(self, in_channels=1, out_channels=1, norm=True):
+        super(IAN, self).__init__(in_channels, out_channels, norm)
+
+
+class ANSN(BaseNet):
+    def __init__(self, in_channels=1, out_channels=1, norm=True):
+        super(ANSN, self).__init__(in_channels, out_channels, norm)
+        self.outc = OutConv(32, out_channels, act=False)
+
+
+class FuseNet(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1, norm=False):
+        super(FuseNet, self).__init__()
+        self.inc = AttentiveDoubleConv(in_channels, 32, norm=norm, leaky=False)
+        self.down1 = AttentiveDown(32, 64, norm=norm, leaky=False)
+        self.down2 = AttentiveDown(64, 64, norm=norm, leaky=False)
+        self.up1 = AttentiveUp(128, 32, bilinear=True, norm=norm, leaky=False)
+        self.up2 = AttentiveUp(64, 32, bilinear=True, norm=norm, leaky=False)
+        self.outc = OutConv(32, out_channels)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x = self.up1(x3, x2)
+        x = self.up2(x, x1)
+        logits = self.outc(x)
+        return logits
+
+
+if __name__ == '__main__':
+    for key in FuseNet(4, 2).state_dict().keys():
+        print(key)