Add files via upload

Author: FanChiMao

generate_patches.py

@@ -0,0 +1,64 @@
+from glob import glob
+from tqdm import tqdm
+import numpy as np
+import os
+from natsort import natsorted
+import cv2
+from joblib import Parallel, delayed
+import argparse
+
+parser = argparse.ArgumentParser(description='Generate patches from Full Resolution images')
+parser.add_argument('--src_dir', default='C:/Users/Lab722 BX/Desktop/IO-Haze/train', type=str, help='Directory for full resolution images')
+parser.add_argument('--tar_dir', default='G:/IO_Haze/train',type=str, help='Directory for image patches')
+parser.add_argument('--ps', default=256, type=int, help='Image Patch Size')
+parser.add_argument('--num_patches', default=200, type=int, help='Number of patches per image')
+parser.add_argument('--num_cores', default=6, type=int, help='Number of CPU Cores')
+
+args = parser.parse_args()
+
+src = args.src_dir
+tar = args.tar_dir
+PS = args.ps
+NUM_PATCHES = args.num_patches
+NUM_CORES = args.num_cores
+
+noisy_patchDir = os.path.join(tar, 'input')
+clean_patchDir = os.path.join(tar, 'target')
+
+if os.path.exists(tar):
+    os.system("rm -r {}".format(tar))
+
+os.makedirs(noisy_patchDir)
+os.makedirs(clean_patchDir)
+
+#get sorted folders
+files = natsorted(glob(os.path.join(src, '*', '*.JPG')))
+
+noisy_files, clean_files = [], []
+for file_ in files:
+    filename = os.path.split(file_)[-1]
+    if 'GT' in filename:
+        clean_files.append(file_)
+    if 'hazy' in filename:
+        noisy_files.append(file_)
+    #if 'gt' in file_:
+    #    clean_files.append(file_)
+    #if 'data' in file_:
+    #    noisy_files.append(file_)
+def save_files(i):
+    noisy_file, clean_file = noisy_files[i], clean_files[i]
+    noisy_img = cv2.imread(noisy_file)
+    clean_img = cv2.imread(clean_file)
+
+    H = noisy_img.shape[0]
+    W = noisy_img.shape[1]
+    for j in range(NUM_PATCHES):
+        rr = np.random.randint(0, H - PS)
+        cc = np.random.randint(0, W - PS)
+        noisy_patch = noisy_img[rr:rr + PS, cc:cc + PS, :]
+        clean_patch = clean_img[rr:rr + PS, cc:cc + PS, :]
+
+        cv2.imwrite(os.path.join(noisy_patchDir, '{}_{}.png'.format(i+1, j+1)), noisy_patch)
+        cv2.imwrite(os.path.join(clean_patchDir, '{}_{}.png'.format(i+1, j+1)), clean_patch)
+
+Parallel(n_jobs=NUM_CORES)(delayed(save_files)(i) for i in tqdm(range(len(noisy_files))))

losses.py

@@ -0,0 +1,167 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+import numpy as np
+from typing import Tuple
+"""
+Reference from:　https://github.com/swz30/MPRNet/blob/main/Denoising/generate_patches_SIDD.py
+"""
+def gaussian(window_size, sigma):
+    def gauss_fcn(x):
+        return -(x - window_size // 2)**2 / float(2 * sigma**2)
+    gauss = torch.stack(
+        [torch.exp(torch.tensor(gauss_fcn(x))) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+def get_gaussian_kernel(ksize: int, sigma: float) -> torch.Tensor:
+    if not isinstance(ksize, int) or ksize % 2 == 0 or ksize <= 0:
+        raise TypeError("ksize must be an odd positive integer. Got {}"
+                        .format(ksize))
+    window_1d: torch.Tensor = gaussian(ksize, sigma)
+    return window_1d
+
+def get_gaussian_kernel2d(ksize: Tuple[int, int],
+                          sigma: Tuple[float, float]) -> torch.Tensor:
+    if not isinstance(ksize, tuple) or len(ksize) != 2:
+        raise TypeError("ksize must be a tuple of length two. Got {}"
+                        .format(ksize))
+    if not isinstance(sigma, tuple) or len(sigma) != 2:
+        raise TypeError("sigma must be a tuple of length two. Got {}"
+                        .format(sigma))
+    ksize_x, ksize_y = ksize
+    sigma_x, sigma_y = sigma
+    kernel_x: torch.Tensor = get_gaussian_kernel(ksize_x, sigma_x)
+    kernel_y: torch.Tensor = get_gaussian_kernel(ksize_y, sigma_y)
+    kernel_2d: torch.Tensor = torch.matmul(
+        kernel_x.unsqueeze(-1), kernel_y.unsqueeze(-1).t())
+    return kernel_2d
+
+
+class PSNRLoss(nn.Module):
+    """
+    reference from: https://github.com/megvii-model/HINet/blob/main/basicsr/models/losses/losses.py
+    """
+    def __init__(self, loss_weight=1.0, reduction='mean', toY=False):
+        super(PSNRLoss, self).__init__()
+        assert reduction == 'mean'
+        self.loss_weight = loss_weight
+        self.scale = 10 / np.log(10)
+        self.toY = toY
+        self.coef = torch.tensor([65.481, 128.553, 24.966]).reshape(1, 3, 1, 1)
+        self.first = True
+
+    def forward(self, pred, target):
+        assert len(pred.size()) == 4
+        if self.toY:
+            if self.first:
+                self.coef = self.coef.to(pred.device)
+                self.first = False
+
+            pred = (pred * self.coef).sum(dim=1).unsqueeze(dim=1) + 16.
+            target = (target * self.coef).sum(dim=1).unsqueeze(dim=1) + 16.
+
+            pred, target = pred / 255., target / 255.
+            pass
+        assert len(pred.size()) == 4
+        loss = -(self.loss_weight * self.scale * torch.log(((pred - target) ** 2).mean(dim=(1, 2, 3)) + 1e-8).mean())
+        return loss
+
+
+class SSIMLoss(nn.Module):
+    def __init__(self, window_size: int = 11, reduction: str = 'mean', max_val: float = 1.0) -> None:
+        super(SSIMLoss, self).__init__()
+        self.window_size: int = window_size
+        self.max_val: float = max_val
+        self.reduction: str = reduction
+
+        self.window: torch.Tensor = get_gaussian_kernel2d(
+            (window_size, window_size), (1.5, 1.5))
+        self.padding: int = self.compute_zero_padding(window_size)
+
+        self.C1: float = (0.01 * self.max_val) ** 2
+        self.C2: float = (0.03 * self.max_val) ** 2
+
+    @staticmethod
+    def compute_zero_padding(kernel_size: int) -> int:
+        """Computes zero padding."""
+        return (kernel_size - 1) // 2
+
+    def filter2D(
+            self,
+            input: torch.Tensor,
+            kernel: torch.Tensor,
+            channel: int) -> torch.Tensor:
+        return F.conv2d(input, kernel, padding=self.padding, groups=channel)
+
+    def forward(self, img1: torch.Tensor, img2: torch.Tensor) -> torch.Tensor:
+        # prepare kernel
+        b, c, h, w = img1.shape
+        tmp_kernel: torch.Tensor = self.window.to(img1.device).to(img1.dtype)
+        kernel: torch.Tensor = tmp_kernel.repeat(c, 1, 1, 1)
+
+        # compute local mean per channel
+        mu1: torch.Tensor = self.filter2D(img1, kernel, c)
+        mu2: torch.Tensor = self.filter2D(img2, kernel, c)
+
+        mu1_sq = mu1.pow(2)
+        mu2_sq = mu2.pow(2)
+        mu1_mu2 = mu1 * mu2
+
+        # compute local sigma per channel
+        sigma1_sq = self.filter2D(img1 * img1, kernel, c) - mu1_sq
+        sigma2_sq = self.filter2D(img2 * img2, kernel, c) - mu2_sq
+        sigma12 = self.filter2D(img1 * img2, kernel, c) - mu1_mu2
+
+        ssim_map = ((2 * mu1_mu2 + self.C1) * (2 * sigma12 + self.C2)) / \
+            ((mu1_sq + mu2_sq + self.C1) * (sigma1_sq + sigma2_sq + self.C2))
+
+        loss = torch.clamp(1. - ssim_map, min=0, max=1) / 2.
+
+        if self.reduction == 'mean':
+            loss = torch.mean(loss)
+        elif self.reduction == 'sum':
+            loss = torch.sum(loss)
+        elif self.reduction == 'none':
+            pass
+        return loss
+# ------------------------------------------------------------------------------
+
+class CharbonnierLoss(nn.Module):
+    """Charbonnier Loss (L1)"""
+
+    def __init__(self, eps=1e-3):
+        super(CharbonnierLoss, self).__init__()
+        self.eps = eps
+
+    def forward(self, x, y):
+        diff = x - y
+        # loss = torch.sum(torch.sqrt(diff * diff + self.eps))
+        loss = torch.mean(torch.sqrt((diff * diff) + (self.eps*self.eps)))
+        return loss
+
+class EdgeLoss(nn.Module):
+    def __init__(self):
+        super(EdgeLoss, self).__init__()
+        k = torch.Tensor([[.05, .25, .4, .25, .05]])
+        self.kernel = torch.matmul(k.t(),k).unsqueeze(0).repeat(3,1,1,1)
+        if torch.cuda.is_available():
+            self.kernel = self.kernel.cuda()
+        self.loss = CharbonnierLoss()
+
+    def conv_gauss(self, img):
+        n_channels, _, kw, kh = self.kernel.shape
+        img = F.pad(img, (kw//2, kh//2, kw//2, kh//2), mode='replicate')
+        return F.conv2d(img, self.kernel, groups=n_channels)
+
+    def laplacian_kernel(self, current):
+        filtered    = self.conv_gauss(current)    # filter
+        down        = filtered[:,:,::2,::2]               # downsample
+        new_filter  = torch.zeros_like(filtered)
+        new_filter[:, :, ::2, ::2] = down*4                  # upsample
+        filtered    = self.conv_gauss(new_filter) # filter
+        diff = current - filtered
+        return diff
+
+    def forward(self, x, y):
+        loss = self.loss(self.laplacian_kernel(x), self.laplacian_kernel(y))
+        return loss