Upload 7 files

README.md

@@ -1,101 +1,49 @@
----
-title: ZeroIG Low-Light Enhancement
-emoji: 🌟
-colorFrom: blue
-colorTo: purple
-sdk: gradio
-sdk_version: 4.44.0
-app_file: app.py
-pinned: false
-license: mit
----
-
-# ZeroIG: Zero-Shot Illumination-Guided Joint Denoising and Adaptive Enhancement
-
-🎉 **CVPR 2024** | Zero-shot low-light image enhancement without training data
-
-## 🚀 Quick Start
-
-Upload a low-light image and get an enhanced version in seconds! No training required.
-
-## 📖 About
-
-This space implements **ZeroIG**, a novel zero-shot method for jointly denoising and enhancing low-light images. The method is completely independent of training data and noise distribution.
-
-### ✨ Key Features
-
-- **Zero-shot**: No training data required
-- **Joint processing**: Simultaneous denoising and enhancement
-- **Illumination-guided**: Smart adaptive enhancement
-- **Prevents artifacts**: Avoids over-enhancement and localized overexposure
-- **Real-time**: Fast processing for practical use
-
-### 🔬 How it Works
-
-1. **Illumination Estimation**: Extracts near-authentic illumination from the input
-2. **Adaptive Enhancement**: Applies different enhancement levels based on pixel intensity  
-3. **Joint Denoising**: Removes noise while preserving image details
-4. **Artifact Prevention**: Prevents common enhancement artifacts
-
-## 📊 Performance
-
-ZeroIG outperforms state-of-the-art methods on standard benchmarks while requiring no training data.
-
-## 🎯 Use Cases
-
-- **Photography**: Rescue underexposed photos
-- **Security**: Enhance surveillance footage  
-- **Mobile**: Real-time camera enhancement
-- **Medical**: Improve low-light medical imaging
-- **Astronomy**: Enhance night sky photography
-
-## 🖼️ Supported Formats
-
-- JPEG, PNG, TIFF, BMP
-- RGB color images
-- Various resolutions (optimized for typical photo sizes)
-
-## ⚡ Tips for Best Results
-
-- Works best with real low-light photos (not artificially darkened)
-- Indoor and outdoor scenes both supported
-- Processing time varies with image size (typically 10-30 seconds)
-
-## 📚 Citation
-
-If you use this work, please cite:
-
-```bibtex
-@inproceedings{shi2024zero,
-    title={ZERO-IG: Zero-Shot Illumination-Guided Joint Denoising and Adaptive Enhancement for Low-Light Images},
-    author={Shi, Yiqi and Liu, Duo and Zhang, Liguo and Tian, Ye and Xia, Xuezhi and Fu, Xiaojing},
-    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
-    pages={3015--3024},
-    year={2024}
-}
-```
-
-## 🔗 Links
-
-- 📄 [Paper](https://openaccess.thecvf.com/content/CVPR2024/papers/Shi_ZERO-IG_Zero-Shot_Illumination-Guided_Joint_Denoising_and_Adaptive_Enhancement_for_Low-Light_CVPR_2024_paper.pdf)
-- 💻 [Code](https://github.com/Doyle59217/ZeroIG)
-- 📊 [Supplement](https://openaccess.thecvf.com/content/CVPR2024/supplemental/Shi_ZERO-IG_Zero-Shot_Illumination-Guided_CVPR_2024_supplemental.pdf)
-
-## 🛠️ Technical Details
-
-- **Framework**: PyTorch
-- **CUDA**: Supported for GPU acceleration
-- **Memory**: Optimized for various image sizes
-- **Dependencies**: See requirements.txt
-
-## 👥 Authors
-
-Yiqi Shi, Duo Liu, Liguo Zhang, Ye Tian, Xuezhi Xia, Xiaojing Fu
-
-## 📄 License
-
-MIT License - see LICENSE file for details
-
----
-
-*Built with ❤️ using Gradio and Hugging Face Spaces*
+# ZERO-IG
+
+### Zero-Shot Illumination-Guided Joint Denoising and Adaptive Enhancement for Low-Light Images [cvpr2024]
+
+By Yiqi Shi, Duo Liu, LiguoZhang,Ye Tian, Xuezhi Xia, Xiaojing Fu
+
+
+#[[Paper]](https://openaccess.thecvf.com/content/CVPR2024/papers/Shi_ZERO-IG_Zero-Shot_Illumination-Guided_Joint_Denoising_and_Adaptive_Enhancement_for_Low-Light_CVPR_2024_paper.pdf)   [[Supplement Material]](https://openaccess.thecvf.com/content/CVPR2024/supplemental/Shi_ZERO-IG_Zero-Shot_Illumination-Guided_CVPR_2024_supplemental.pdf)
+
+# Zero-IG Framework
+
+<img src="Figs/Fig3.png" width="900px"/> 
+<p style="text-align:justify">Note that the provided model in this code are not the model for generating results reported in the paper.
+
+## Model Training Configuration
+* To train a new model, specify the dataset path in "train.py" and execute it. The trained model will be stored in the 'weights' folder, while intermediate visualization outputs will be saved in the 'results' folder.
+* We have provided some model parameters, but we recommend training with a single image for better result.
+
+## Requirements
+* Python 3.7
+* PyTorch 1.13.0
+* CUDA 11.7
+* Torchvision 0.14.1
+
+## Testing
+* Ensure the data is prepared and placed in the designated folder.
+* Select the appropriate model for testing, which could be a model trained by yourself.
+* Execute "test.py" to perform the testing.
+
+## [VILNC Dataset](https://pan.baidu.com/s/1-Uw78IxlVAVY_hqRRS9BGg?pwd=4e5c )
+
+The Varied Indoor Luminance & Nightscapes Collection (VILNC Dataset) is a meticulously curated assembly of 500 real-world low-light images, captured with the precision of a Canon EOS 550D camera. This dataset is segmented into two main environments, comprising 460 indoor scenes and 40 outdoor landscapes. Within the indoor category, each scene is represented through a trio of images, each depicting a distinct level of dim luminance, alongside a corresponding reference image captured under normal lighting conditions. For the outdoor scenes, the dataset includes low-light photographs, each paired with its respective normal light reference image, providing a comprehensive resource for analyzing and enhancing low-light imaging techniques.
+
+<img src="Figs/Dataset.png" width="900px"/> 
+<p style="text-align:justify">
+
+
+
+## Citation
+```bibtex
+@inproceedings{shi2024zero,
+  title={ZERO-IG: Zero-Shot Illumination-Guided Joint Denoising and Adaptive Enhancement for Low-Light Images},
+  author={Shi, Yiqi and Liu, Duo and Zhang, Liguo and Tian, Ye and Xia, Xuezhi and Fu, Xiaojing},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={3015--3024},
+  year={2024}
+}
+```
+

loss.py

@@ -0,0 +1,307 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import scipy.stats as st
+from utils import  pair_downsampler,calculate_local_variance,LocalMean
+
+EPS = 1e-9
+PI = 22.0 / 7.0
+
+
+class LossFunction(nn.Module):
+    def __init__(self):
+        super(LossFunction, self).__init__()
+        self._l2_loss = nn.MSELoss()
+        self._l1_loss = nn.L1Loss()
+        self.smooth_loss = SmoothLoss()
+        self.texture_difference=TextureDifference()
+        self.local_mean=LocalMean(patch_size=5)
+        self.L_TV_loss=L_TV()
+
+
+    def forward(self,input,L_pred1,L_pred2,L2,s2,s21,s22,H2,H11,H12,H13,s13,H14,s14,H3,s3,H3_pred,H4_pred,L_pred1_L_pred2_diff,H3_denoised1_H3_denoised2_diff,H2_blur,H3_blur):
+        eps = 1e-9
+        input = input + eps
+
+        input_Y = L2.detach()[:, 2, :, :] * 0.299 + L2.detach()[:, 1, :, :] * 0.587 + L2.detach()[:, 0, :, :] * 0.144
+        input_Y_mean = torch.mean(input_Y, dim=(1, 2))
+        enhancement_factor = 0.5/ (input_Y_mean + eps)
+        enhancement_factor = enhancement_factor.unsqueeze(1).unsqueeze(2).unsqueeze(3)
+        enhancement_factor = torch.clamp(enhancement_factor, 1, 25)
+        adjustment_ratio = torch.pow(0.7, -enhancement_factor) / enhancement_factor
+        adjustment_ratio = adjustment_ratio.repeat(1, 3, 1, 1)
+        normalized_low_light_layer  = L2.detach() / s2
+        normalized_low_light_layer = torch.clamp(normalized_low_light_layer, eps, 0.8)
+        enhanced_brightness=torch.pow(L2.detach()*enhancement_factor, enhancement_factor)
+        clamped_enhanced_brightness = torch.clamp(enhanced_brightness * adjustment_ratio, eps, 1)
+        clamped_adjusted_low_light  = torch.clamp(L2.detach() *  enhancement_factor,eps,1)
+        loss = 0
+        #Enhance_loss
+        loss += self._l2_loss(s2, clamped_enhanced_brightness) *700
+        loss += self._l2_loss(normalized_low_light_layer, clamped_adjusted_low_light) *1000
+        loss += self.smooth_loss(L2.detach(), s2) *5
+        loss += self.L_TV_loss(s2)*1600
+        #Loss_res_1
+        L11, L12 = pair_downsampler(input)
+        loss += self._l2_loss(L11, L_pred2) * 1000
+        loss += self._l2_loss(L12, L_pred1) * 1000
+        denoised1, denoised2 = pair_downsampler(L2)
+        loss += self._l2_loss(L_pred1, denoised1) * 1000
+        loss += self._l2_loss(L_pred2, denoised2) * 1000
+        # Loss_res_2
+        loss += self._l2_loss(H3_pred, torch.cat([H12.detach(), s22.detach()], 1)) * 1000
+        loss += self._l2_loss(H4_pred, torch.cat([H11.detach(), s21.detach()], 1)) * 1000
+        H3_denoised1, H3_denoised2 = pair_downsampler(H3)
+        loss += self._l2_loss(H3_pred[:, 0:3, :, :], H3_denoised1) * 1000
+        loss += self._l2_loss(H4_pred[:, 0:3, :, :], H3_denoised2) * 1000
+        #Loss_color
+        loss += self._l2_loss(H2_blur.detach(), H3_blur) * 10000
+        #Loss_ill
+        loss += self._l2_loss(s2.detach(), s3) * 1000
+        #Loss_cons
+        local_mean1 = self.local_mean(H3_denoised1)
+        local_mean2 = self.local_mean(H3_denoised2)
+        weighted_diff1 = (1 - H3_denoised1_H3_denoised2_diff) * local_mean1+H3_denoised1*H3_denoised1_H3_denoised2_diff
+        weighted_diff2 = (1 - H3_denoised1_H3_denoised2_diff) * local_mean2+H3_denoised1*H3_denoised1_H3_denoised2_diff
+        loss += self._l2_loss(H3_denoised1,weighted_diff1)* 10000
+        loss += self._l2_loss(H3_denoised2, weighted_diff2)* 10000
+        #Loss_Var
+        noise_std = calculate_local_variance(H3 - H2)
+        H2_var = calculate_local_variance(H2)
+        loss += self._l2_loss(H2_var, noise_std) * 1000
+        return loss
+
+def local_mean(self, image):
+    padding = self.patch_size // 2
+    image = F.pad(image, (padding, padding, padding, padding), mode='reflect')
+    patches = image.unfold(2, self.patch_size, 1).unfold(3, self.patch_size, 1)
+    return patches.mean(dim=(4, 5))
+
+def gauss_kernel(kernlen=21, nsig=3, channels=1):
+    interval = (2 * nsig + 1.) / (kernlen)
+    x = np.linspace(-nsig - interval / 2., nsig + interval / 2., kernlen + 1)
+    kern1d = np.diff(st.norm.cdf(x))
+    kernel_raw = np.sqrt(np.outer(kern1d, kern1d))
+    kernel = kernel_raw / kernel_raw.sum()
+    out_filter = np.array(kernel, dtype=np.float32)
+    out_filter = out_filter.reshape((kernlen, kernlen, 1, 1))
+    out_filter = np.repeat(out_filter, channels, axis=2)
+
+    return out_filter
+
+
+class TextureDifference(nn.Module):
+    def __init__(self, patch_size=5, constant_C=1e-5,threshold=0.975):
+        super(TextureDifference, self).__init__()
+        self.patch_size = patch_size
+        self.constant_C = constant_C
+        self.threshold = threshold
+
+    def forward(self, image1, image2):
+        # Convert RGB images to grayscale
+        image1 = self.rgb_to_gray(image1)
+        image2 = self.rgb_to_gray(image2)
+
+        stddev1 = self.local_stddev(image1)
+        stddev2 = self.local_stddev(image2)
+        numerator = 2 * stddev1 * stddev2
+        denominator = stddev1 ** 2 + stddev2 ** 2 + self.constant_C
+        diff = numerator / denominator
+
+        # Apply threshold to diff tensor
+        binary_diff = torch.where(diff > self.threshold, torch.tensor(1.0, device=diff.device),
+                                  torch.tensor(0.0, device=diff.device))
+
+        return binary_diff
+
+    def local_stddev(self, image):
+        padding = self.patch_size // 2
+        image = F.pad(image, (padding, padding, padding, padding), mode='reflect')
+        patches = image.unfold(2, self.patch_size, 1).unfold(3, self.patch_size, 1)
+        mean = patches.mean(dim=(4, 5), keepdim=True)
+        squared_diff = (patches - mean) ** 2
+        local_variance = squared_diff.mean(dim=(4, 5))
+        local_stddev = torch.sqrt(local_variance+1e-9)
+        return local_stddev
+
+    def rgb_to_gray(self, image):
+        # Convert RGB image to grayscale using the luminance formula
+        gray_image =  0.144 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.299 * image[:, 2, :, :]
+        return gray_image.unsqueeze(1)  # Add a channel dimension for compatibility
+
+
+class L_TV(nn.Module):
+    def __init__(self,TVLoss_weight=1):
+        super(L_TV,self).__init__()
+        self.TVLoss_weight = TVLoss_weight
+
+    def forward(self,x):
+        batch_size = x.size()[0]
+        h_x = x.size()[2]
+        w_x = x.size()[3]
+        count_h =  (x.size()[2]-1) * x.size()[3]
+        count_w = x.size()[2] * (x.size()[3] - 1)
+        h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum()
+        w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum()
+        return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size
+
+class Blur(nn.Module):
+    def __init__(self, nc):
+        super(Blur, self).__init__()
+        self.nc = nc
+        kernel = gauss_kernel(kernlen=21, nsig=3, channels=self.nc)
+        kernel = torch.from_numpy(kernel).permute(2, 3, 0, 1).cuda()
+        self.weight = nn.Parameter(data=kernel, requires_grad=False).cuda()
+
+    def forward(self, x):
+        if x.size(1) != self.nc:
+            raise RuntimeError(
+                "The channel of input [%d] does not match the preset channel [%d]" % (x.size(1), self.nc))
+
+        x = F.conv2d(x, self.weight, stride=1, padding=10, groups=self.nc)
+        return x
+
+
+
+
+class SmoothLoss(nn.Module):
+    def __init__(self):
+        super(SmoothLoss, self).__init__()
+        self.sigma = 10
+
+    def rgb2yCbCr(self, input_im):
+
+        im_flat = input_im.contiguous().view(-1, 3).float()
+        # [w,h,3] => [w*h,3]
+        mat = torch.Tensor([[0.257, -0.148, 0.439], [0.564, -0.291, -0.368], [0.098, 0.439, -0.071]]).cuda()
+        # [3,3]
+        bias = torch.Tensor([16.0 / 255.0, 128.0 / 255.0, 128.0 / 255.0]).cuda()
+        # [1,3]
+        temp = im_flat.mm(mat) + bias
+        # [w*h,3]*[3,3]+[1,3] => [w*h,3]
+        out = temp.view(input_im.shape[0], 3, input_im.shape[2], input_im.shape[3])
+        return out
+
+    # output: output      input:input
+    def forward(self, input, output):
+
+
+        self.output = output
+        self.input = self.rgb2yCbCr(input)
+        sigma_color = -1.0 / (2 * self.sigma * self.sigma)
+        w1 = torch.exp(torch.sum(torch.pow(self.input[:, :, 1:, :] - self.input[:, :, :-1, :], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w2 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-1, :] - self.input[:, :, 1:, :], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w3 = torch.exp(torch.sum(torch.pow(self.input[:, :, :, 1:] - self.input[:, :, :, :-1], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w4 = torch.exp(torch.sum(torch.pow(self.input[:, :, :, :-1] - self.input[:, :, :, 1:], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w5 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-1, :-1] - self.input[:, :, 1:, 1:], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w6 = torch.exp(torch.sum(torch.pow(self.input[:, :, 1:, 1:] - self.input[:, :, :-1, :-1], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w7 = torch.exp(torch.sum(torch.pow(self.input[:, :, 1:, :-1] - self.input[:, :, :-1, 1:], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w8 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-1, 1:] - self.input[:, :, 1:, :-1], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w9 = torch.exp(torch.sum(torch.pow(self.input[:, :, 2:, :] - self.input[:, :, :-2, :], 2), dim=1,
+                                 keepdim=True) * sigma_color)
+        w10 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-2, :] - self.input[:, :, 2:, :], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w11 = torch.exp(torch.sum(torch.pow(self.input[:, :, :, 2:] - self.input[:, :, :, :-2], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w12 = torch.exp(torch.sum(torch.pow(self.input[:, :, :, :-2] - self.input[:, :, :, 2:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w13 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-2, :-1] - self.input[:, :, 2:, 1:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w14 = torch.exp(torch.sum(torch.pow(self.input[:, :, 2:, 1:] - self.input[:, :, :-2, :-1], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w15 = torch.exp(torch.sum(torch.pow(self.input[:, :, 2:, :-1] - self.input[:, :, :-2, 1:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w16 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-2, 1:] - self.input[:, :, 2:, :-1], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w17 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-1, :-2] - self.input[:, :, 1:, 2:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w18 = torch.exp(torch.sum(torch.pow(self.input[:, :, 1:, 2:] - self.input[:, :, :-1, :-2], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w19 = torch.exp(torch.sum(torch.pow(self.input[:, :, 1:, :-2] - self.input[:, :, :-1, 2:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w20 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-1, 2:] - self.input[:, :, 1:, :-2], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w21 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-2, :-2] - self.input[:, :, 2:, 2:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w22 = torch.exp(torch.sum(torch.pow(self.input[:, :, 2:, 2:] - self.input[:, :, :-2, :-2], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w23 = torch.exp(torch.sum(torch.pow(self.input[:, :, 2:, :-2] - self.input[:, :, :-2, 2:], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        w24 = torch.exp(torch.sum(torch.pow(self.input[:, :, :-2, 2:] - self.input[:, :, 2:, :-2], 2), dim=1,
+                                  keepdim=True) * sigma_color)
+        p = 1.0
+
+        pixel_grad1 = w1 * torch.norm((self.output[:, :, 1:, :] - self.output[:, :, :-1, :]), p, dim=1, keepdim=True)
+        pixel_grad2 = w2 * torch.norm((self.output[:, :, :-1, :] - self.output[:, :, 1:, :]), p, dim=1, keepdim=True)
+        pixel_grad3 = w3 * torch.norm((self.output[:, :, :, 1:] - self.output[:, :, :, :-1]), p, dim=1, keepdim=True)
+        pixel_grad4 = w4 * torch.norm((self.output[:, :, :, :-1] - self.output[:, :, :, 1:]), p, dim=1, keepdim=True)
+        pixel_grad5 = w5 * torch.norm((self.output[:, :, :-1, :-1] - self.output[:, :, 1:, 1:]), p, dim=1, keepdim=True)
+        pixel_grad6 = w6 * torch.norm((self.output[:, :, 1:, 1:] - self.output[:, :, :-1, :-1]), p, dim=1, keepdim=True)
+        pixel_grad7 = w7 * torch.norm((self.output[:, :, 1:, :-1] - self.output[:, :, :-1, 1:]), p, dim=1, keepdim=True)
+        pixel_grad8 = w8 * torch.norm((self.output[:, :, :-1, 1:] - self.output[:, :, 1:, :-1]), p, dim=1, keepdim=True)
+        pixel_grad9 = w9 * torch.norm((self.output[:, :, 2:, :] - self.output[:, :, :-2, :]), p, dim=1, keepdim=True)
+        pixel_grad10 = w10 * torch.norm((self.output[:, :, :-2, :] - self.output[:, :, 2:, :]), p, dim=1, keepdim=True)
+        pixel_grad11 = w11 * torch.norm((self.output[:, :, :, 2:] - self.output[:, :, :, :-2]), p, dim=1, keepdim=True)
+        pixel_grad12 = w12 * torch.norm((self.output[:, :, :, :-2] - self.output[:, :, :, 2:]), p, dim=1, keepdim=True)
+        pixel_grad13 = w13 * torch.norm((self.output[:, :, :-2, :-1] - self.output[:, :, 2:, 1:]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad14 = w14 * torch.norm((self.output[:, :, 2:, 1:] - self.output[:, :, :-2, :-1]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad15 = w15 * torch.norm((self.output[:, :, 2:, :-1] - self.output[:, :, :-2, 1:]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad16 = w16 * torch.norm((self.output[:, :, :-2, 1:] - self.output[:, :, 2:, :-1]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad17 = w17 * torch.norm((self.output[:, :, :-1, :-2] - self.output[:, :, 1:, 2:]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad18 = w18 * torch.norm((self.output[:, :, 1:, 2:] - self.output[:, :, :-1, :-2]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad19 = w19 * torch.norm((self.output[:, :, 1:, :-2] - self.output[:, :, :-1, 2:]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad20 = w20 * torch.norm((self.output[:, :, :-1, 2:] - self.output[:, :, 1:, :-2]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad21 = w21 * torch.norm((self.output[:, :, :-2, :-2] - self.output[:, :, 2:, 2:]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad22 = w22 * torch.norm((self.output[:, :, 2:, 2:] - self.output[:, :, :-2, :-2]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad23 = w23 * torch.norm((self.output[:, :, 2:, :-2] - self.output[:, :, :-2, 2:]), p, dim=1,
+                                        keepdim=True)
+        pixel_grad24 = w24 * torch.norm((self.output[:, :, :-2, 2:] - self.output[:, :, 2:, :-2]), p, dim=1,
+                                        keepdim=True)
+
+        ReguTerm1 = torch.mean(pixel_grad1) \
+                    + torch.mean(pixel_grad2) \
+                    + torch.mean(pixel_grad3) \
+                    + torch.mean(pixel_grad4) \
+                    + torch.mean(pixel_grad5) \
+                    + torch.mean(pixel_grad6) \
+                    + torch.mean(pixel_grad7) \
+                    + torch.mean(pixel_grad8) \
+                    + torch.mean(pixel_grad9) \
+                    + torch.mean(pixel_grad10) \
+                    + torch.mean(pixel_grad11) \
+                    + torch.mean(pixel_grad12) \
+                    + torch.mean(pixel_grad13) \
+                    + torch.mean(pixel_grad14) \
+                    + torch.mean(pixel_grad15) \
+                    + torch.mean(pixel_grad16) \
+                    + torch.mean(pixel_grad17) \
+                    + torch.mean(pixel_grad18) \
+                    + torch.mean(pixel_grad19) \
+                    + torch.mean(pixel_grad20) \
+                    + torch.mean(pixel_grad21) \
+                    + torch.mean(pixel_grad22) \
+                    + torch.mean(pixel_grad23) \
+                    + torch.mean(pixel_grad24)
+
+        total_term = ReguTerm1
+        return total_term
+

model.py

@@ -1,271 +1,207 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def pair_downsampler(img):
-    # img has shape B C H W
-    c = img.shape[1]
-    filter1 = torch.FloatTensor([[[[0, 0.5], [0.5, 0]]]]).to(img.device)
-    filter1 = filter1.repeat(c, 1, 1, 1)
-    filter2 = torch.FloatTensor([[[[0.5, 0], [0, 0.5]]]]).to(img.device)
-    filter2 = filter2.repeat(c, 1, 1, 1)
-    output1 = torch.nn.functional.conv2d(img, filter1, stride=2, groups=c)
-    output2 = torch.nn.functional.conv2d(img, filter2, stride=2, groups=c)
-    return output1, output2
-
-
-def gauss_cdf(x):
-    return 0.5*(1+torch.erf(x/torch.sqrt(torch.tensor(2.))))
-
-
-def gauss_kernel(kernlen=21, nsig=3, channels=1):
-    interval = (2*nsig+1.)/(kernlen)
-    x = torch.linspace(-nsig-interval/2., nsig+interval/2., kernlen+1).to('cuda' if torch.cuda.is_available() else 'cpu')
-    kern1d = torch.diff(gauss_cdf(x))
-    kernel_raw = torch.sqrt(torch.outer(kern1d, kern1d))
-    kernel = kernel_raw/torch.sum(kernel_raw)
-    out_filter = kernel.view(1, 1, kernlen, kernlen)
-    out_filter = out_filter.repeat(channels, 1, 1, 1)
-    return out_filter
-
-
-def blur(x):
-    device = x.device
-    kernel_size = 21
-    padding = kernel_size // 2
-    kernel_var = gauss_kernel(kernel_size, 1, x.size(1)).to(device)
-    x_padded = torch.nn.functional.pad(x, (padding, padding, padding, padding), mode='reflect')
-    return torch.nn.functional.conv2d(x_padded, kernel_var, padding=0, groups=x.size(1))
-
-
-class TextureDifference(nn.Module):
-    def __init__(self, patch_size=5, constant_C=1e-5, threshold=0.975):
-        super(TextureDifference, self).__init__()
-        self.patch_size = patch_size
-        self.constant_C = constant_C
-        self.threshold = threshold
-
-    def forward(self, image1, image2):
-        # Convert RGB images to grayscale
-        image1 = self.rgb_to_gray(image1)
-        image2 = self.rgb_to_gray(image2)
-
-        stddev1 = self.local_stddev(image1)
-        stddev2 = self.local_stddev(image2)
-        numerator = 2 * stddev1 * stddev2
-        denominator = stddev1 ** 2 + stddev2 ** 2 + self.constant_C
-        diff = numerator / denominator
-
-        # Apply threshold to diff tensor
-        binary_diff = torch.where(diff > self.threshold, torch.tensor(1.0, device=diff.device),
-                                  torch.tensor(0.0, device=diff.device))
-
-        return binary_diff
-
-    def local_stddev(self, image):
-        padding = self.patch_size // 2
-        image = F.pad(image, (padding, padding, padding, padding), mode='reflect')
-        patches = image.unfold(2, self.patch_size, 1).unfold(3, self.patch_size, 1)
-        mean = patches.mean(dim=(4, 5), keepdim=True)
-        squared_diff = (patches - mean) ** 2
-        local_variance = squared_diff.mean(dim=(4, 5))
-        local_stddev = torch.sqrt(local_variance+1e-9)
-        return local_stddev
-
-    def rgb_to_gray(self, image):
-        # Convert RGB image to grayscale using the luminance formula
-        gray_image = 0.144 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.299 * image[:, 2, :, :]
-        return gray_image.unsqueeze(1)  # Add a channel dimension for compatibility
-
-
-class Denoise_1(nn.Module):
-    def __init__(self, chan_embed=48):
-        super(Denoise_1, self).__init__()
-
-        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
-        self.conv1 = nn.Conv2d(3, chan_embed, 3, padding=1)
-        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding=1)
-        self.conv3 = nn.Conv2d(chan_embed, 3, 1)
-
-    def forward(self, x):
-        x = self.act(self.conv1(x))
-        x = self.act(self.conv2(x))
-        x = self.conv3(x)
-        return x
-
-
-class Denoise_2(nn.Module):
-    def __init__(self, chan_embed=96):
-        super(Denoise_2, self).__init__()
-
-        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
-        self.conv1 = nn.Conv2d(6, chan_embed, 3, padding=1)
-        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding=1)
-        self.conv3 = nn.Conv2d(chan_embed, 6, 1)
-
-    def forward(self, x):
-        x = self.act(self.conv1(x))
-        x = self.act(self.conv2(x))
-        x = self.conv3(x)
-        return x
-
-
-class Enhancer(nn.Module):
-    def __init__(self, layers, channels):
-        super(Enhancer, self).__init__()
-
-        kernel_size = 3
-        dilation = 1
-        padding = int((kernel_size - 1) / 2) * dilation
-
-        self.in_conv = nn.Sequential(
-            nn.Conv2d(in_channels=3, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
-            nn.ReLU()
-        )
-
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
-            nn.BatchNorm2d(channels),
-            nn.ReLU()
-        )
-        self.blocks = nn.ModuleList()
-        for i in range(layers):
-            self.blocks.append(self.conv)
-
-        self.out_conv = nn.Sequential(
-            nn.Conv2d(in_channels=channels, out_channels=3, kernel_size=3, stride=1, padding=1),
-            nn.Sigmoid()
-        )
-
-    def forward(self, input):
-        fea = self.in_conv(input)
-        for conv in self.blocks:
-            fea = fea + conv(fea)
-        fea = self.out_conv(fea)
-        fea = torch.clamp(fea, 0.0001, 1)
-
-        return fea
-
-
-class Network(nn.Module):
-    def __init__(self):
-        super(Network, self).__init__()
-
-        self.enhance = Enhancer(layers=3, channels=64)
-        self.denoise_1 = Denoise_1(chan_embed=48)
-        self.denoise_2 = Denoise_2(chan_embed=48)
-        self.TextureDifference = TextureDifference()
-
-    def enhance_weights_init(self, m):
-        if isinstance(m, nn.Conv2d):
-            m.weight.data.normal_(0.0, 0.02)
-            if m.bias != None:
-                m.bias.data.zero_()
-
-        if isinstance(m, nn.BatchNorm2d):
-            m.weight.data.normal_(1., 0.02)
-
-    def denoise_weights_init(self, m):
-        if isinstance(m, nn.Conv2d):
-            m.weight.data.normal_(0, 0.02)
-            if m.bias != None:
-                m.bias.data.zero_()
-
-        if isinstance(m, nn.BatchNorm2d):
-            m.weight.data.normal_(1., 0.02)
-
-    def forward(self, input):
-        eps = 1e-4
-        input = input + eps
-
-        L11, L12 = pair_downsampler(input)
-        L_pred1 = L11 - self.denoise_1(L11)
-        L_pred2 = L12 - self.denoise_1(L12)
-        L2 = input - self.denoise_1(input)
-        L2 = torch.clamp(L2, eps, 1)
-
-        s2 = self.enhance(L2.detach())
-        s21, s22 = pair_downsampler(s2)
-        H2 = input / s2
-        H2 = torch.clamp(H2, eps, 1)
-
-        H11 = L11 / s21
-        H11 = torch.clamp(H11, eps, 1)
-
-        H12 = L12 / s22
-        H12 = torch.clamp(H12, eps, 1)
-
-        H3_pred = torch.cat([H11, s21], 1).detach() - self.denoise_2(torch.cat([H11, s21], 1))
-        H3_pred = torch.clamp(H3_pred, eps, 1)
-        H13 = H3_pred[:, :3, :, :]
-        s13 = H3_pred[:, 3:, :, :]
-
-        H4_pred = torch.cat([H12, s22], 1).detach() - self.denoise_2(torch.cat([H12, s22], 1))
-        H4_pred = torch.clamp(H4_pred, eps, 1)
-        H14 = H4_pred[:, :3, :, :]
-        s14 = H4_pred[:, 3:, :, :]
-
-        H5_pred = torch.cat([H2, s2], 1).detach() - self.denoise_2(torch.cat([H2, s2], 1))
-        H5_pred = torch.clamp(H5_pred, eps, 1)
-        H3 = H5_pred[:, :3, :, :]
-        s3 = H5_pred[:, 3:, :, :]
-
-        L_pred1_L_pred2_diff = self.TextureDifference(L_pred1, L_pred2)
-        H3_denoised1, H3_denoised2 = pair_downsampler(H3)
-        H3_denoised1_H3_denoised2_diff = self.TextureDifference(H3_denoised1, H3_denoised2)
-
-        H1 = L2 / s2
-        H1 = torch.clamp(H1, 0, 1)
-        H2_blur = blur(H1)
-        H3_blur = blur(H3)
-
-        return L_pred1, L_pred2, L2, s2, s21, s22, H2, H11, H12, H13, s13, H14, s14, H3, s3, H3_pred, H4_pred, L_pred1_L_pred2_diff, H3_denoised1_H3_denoised2_diff, H2_blur, H3_blur
-
-
-class Finetunemodel(nn.Module):
-    def __init__(self, weights):
-        super(Finetunemodel, self).__init__()
-
-        self.enhance = Enhancer(layers=3, channels=64)
-        self.denoise_1 = Denoise_1(chan_embed=48)
-        self.denoise_2 = Denoise_2(chan_embed=48)
-
-        # Try to load weights if file exists
-        if weights and torch.cuda.is_available():
-            device = 'cuda:0'
-        else:
-            device = 'cpu'
-            
-        try:
-            base_weights = torch.load(weights, map_location=device)
-            pretrained_dict = base_weights
-            model_dict = self.state_dict()
-            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
-            model_dict.update(pretrained_dict)
-            self.load_state_dict(model_dict)
-            print(f"Successfully loaded weights from {weights}")
-        except Exception as e:
-            print(f"Warning: Could not load weights from {weights}: {e}")
-            print("Using randomly initialized weights")
-
-    def weights_init(self, m):
-        if isinstance(m, nn.Conv2d):
-            m.weight.data.normal_(0, 0.02)
-            m.bias.data.zero_()
-
-        if isinstance(m, nn.BatchNorm2d):
-            m.weight.data.normal_(1., 0.02)
-
-    def forward(self, input):
-        eps = 1e-4
-        input = input + eps
-        L2 = input - self.denoise_1(input)
-        L2 = torch.clamp(L2, eps, 1)
-        s2 = self.enhance(L2)
-        H2 = input / s2
-        H2 = torch.clamp(H2, eps, 1)
-        H5_pred = torch.cat([H2, s2], 1).detach() - self.denoise_2(torch.cat([H2, s2], 1))
-        H5_pred = torch.clamp(H5_pred, eps, 1)
-        H3 = H5_pred[:, :3, :, :]
-        return H2, H3
+import torch
+import torch.nn as nn
+from loss import LossFunction, TextureDifference
+from utils import blur, pair_downsampler
+
+
+
+class Denoise_1(nn.Module):
+    def __init__(self, chan_embed=48):
+        super(Denoise_1, self).__init__()
+
+        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.conv1 = nn.Conv2d(3, chan_embed, 3, padding=1)
+        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding=1)
+        self.conv3 = nn.Conv2d(chan_embed, 3, 1)
+
+    def forward(self, x):
+        x = self.act(self.conv1(x))
+        x = self.act(self.conv2(x))
+        x = self.conv3(x)
+        return x
+
+
+class Denoise_2(nn.Module):
+    def __init__(self, chan_embed=96):
+        super(Denoise_2, self).__init__()
+
+        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.conv1 = nn.Conv2d(6, chan_embed, 3, padding=1)
+        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding=1)
+        self.conv3 = nn.Conv2d(chan_embed, 6, 1)
+
+    def forward(self, x):
+        x = self.act(self.conv1(x))
+        x = self.act(self.conv2(x))
+        x = self.conv3(x)
+        return x
+
+
+class Enhancer(nn.Module):
+    def __init__(self, layers, channels):
+        super(Enhancer, self).__init__()
+
+        kernel_size = 3
+        dilation = 1
+        padding = int((kernel_size - 1) / 2) * dilation
+
+        self.in_conv = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
+            nn.ReLU()
+        )
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
+            nn.BatchNorm2d(channels),
+            nn.ReLU()
+        )
+        self.blocks = nn.ModuleList()
+        for i in range(layers):
+            self.blocks.append(self.conv)
+
+        self.out_conv = nn.Sequential(
+            nn.Conv2d(in_channels=channels, out_channels=3, kernel_size=3, stride=1, padding=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, input):
+        fea = self.in_conv(input)
+        for conv in self.blocks:
+            fea = fea + conv(fea)
+        fea = self.out_conv(fea)
+        fea = torch.clamp(fea, 0.0001, 1)
+
+        return fea
+
+
+class Network(nn.Module):
+
+    def __init__(self):
+        super(Network, self).__init__()
+
+        self.enhance = Enhancer(layers=3, channels=64)
+        self.denoise_1 = Denoise_1(chan_embed=48)
+        self.denoise_2 = Denoise_2(chan_embed=48)
+        self._l2_loss = nn.MSELoss()
+        self._l1_loss = nn.L1Loss()
+        self._criterion = LossFunction()
+        self.avgpool = nn.AvgPool2d(kernel_size=3, stride=1, padding=1)
+        self.TextureDifference = TextureDifference()
+
+
+    def enhance_weights_init(self, m):
+        if isinstance(m, nn.Conv2d):
+            m.weight.data.normal_(0.0, 0.02)
+            if m.bias != None:
+                m.bias.data.zero_()
+
+        if isinstance(m, nn.BatchNorm2d):
+            m.weight.data.normal_(1., 0.02)
+
+    def denoise_weights_init(self, m):
+        if isinstance(m, nn.Conv2d):
+            m.weight.data.normal_(0, 0.02)
+            if m.bias != None:
+                m.bias.data.zero_()
+
+        if isinstance(m, nn.BatchNorm2d):
+            m.weight.data.normal_(1., 0.02)
+        # if isinstance(m, nn.Conv2d):
+        # nn.init.xavier_uniform(m.weight)
+        # nn.init.constant(m.bias, 0)
+
+    def forward(self, input):
+        eps = 1e-4
+        input = input + eps
+
+        L11, L12 = pair_downsampler(input)
+        L_pred1 = L11 - self.denoise_1(L11)
+        L_pred2 = L12 - self.denoise_1(L12)
+        L2 = input - self.denoise_1(input)
+        L2 = torch.clamp(L2, eps, 1)
+
+        s2 = self.enhance(L2.detach())
+        s21, s22 = pair_downsampler(s2)
+        H2 = input / s2
+        H2 = torch.clamp(H2, eps, 1)
+
+        H11 = L11 / s21
+        H11 = torch.clamp(H11, eps, 1)
+
+        H12 = L12 / s22
+        H12 = torch.clamp(H12, eps, 1)
+
+        H3_pred = torch.cat([H11, s21], 1).detach() - self.denoise_2(torch.cat([H11, s21], 1))
+        H3_pred = torch.clamp(H3_pred, eps, 1)
+        H13 = H3_pred[:, :3, :, :]
+        s13 = H3_pred[:, 3:, :, :]
+
+        H4_pred = torch.cat([H12, s22], 1).detach() - self.denoise_2(torch.cat([H12, s22], 1))
+        H4_pred = torch.clamp(H4_pred, eps, 1)
+        H14 = H4_pred[:, :3, :, :]
+        s14 = H4_pred[:, 3:, :, :]
+
+        H5_pred = torch.cat([H2, s2], 1).detach() - self.denoise_2(torch.cat([H2, s2], 1))
+        H5_pred = torch.clamp(H5_pred, eps, 1)
+        H3 = H5_pred[:, :3, :, :]
+        s3 = H5_pred[:, 3:, :, :]
+
+        L_pred1_L_pred2_diff = self.TextureDifference(L_pred1, L_pred2)
+        H3_denoised1, H3_denoised2 = pair_downsampler(H3)
+        H3_denoised1_H3_denoised2_diff= self.TextureDifference(H3_denoised1, H3_denoised2)
+
+        H1 = L2 / s2
+        H1 = torch.clamp(H1, 0, 1)
+        H2_blur = blur(H1)
+        H3_blur = blur(H3)
+
+        return L_pred1, L_pred2, L2, s2, s21, s22, H2, H11, H12, H13, s13, H14, s14, H3, s3, H3_pred, H4_pred, L_pred1_L_pred2_diff, H3_denoised1_H3_denoised2_diff, H2_blur, H3_blur
+
+    def _loss(self, input):
+        L_pred1, L_pred2, L2, s2, s21, s22, H2, H11, H12, H13, s13, H14, s14, H3, s3, H3_pred, H4_pred, L_pred1_L_pred2_diff, H3_denoised1_H3_denoised2_diff, H2_blur, H3_blur = self(
+            input)
+        loss = 0
+
+        loss += self._criterion(input, L_pred1, L_pred2, L2, s2, s21, s22, H2, H11, H12, H13, s13, H14, s14, H3, s3,
+                                H3_pred, H4_pred, L_pred1_L_pred2_diff, H3_denoised1_H3_denoised2_diff, H2_blur,
+                                H3_blur)
+        return loss
+
+
+class Finetunemodel(nn.Module):
+
+    def __init__(self, weights):
+        super(Finetunemodel, self).__init__()
+
+        self.enhance = Enhancer(layers=3, channels=64)
+        self.denoise_1 = Denoise_1(chan_embed=48)
+        self.denoise_2 = Denoise_2(chan_embed=48)
+
+        base_weights = torch.load(weights, map_location='cuda:0')
+        pretrained_dict = base_weights
+        model_dict = self.state_dict()
+        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+        model_dict.update(pretrained_dict)
+        self.load_state_dict(model_dict)
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Conv2d):
+            m.weight.data.normal_(0, 0.02)
+            m.bias.data.zero_()
+
+        if isinstance(m, nn.BatchNorm2d):
+            m.weight.data.normal_(1., 0.02)
+
+    def forward(self, input):
+        eps = 1e-4
+        input = input + eps
+        L2 = input - self.denoise_1(input)
+        L2 = torch.clamp(L2, eps, 1)
+        s2 = self.enhance(L2)
+        H2 = input / s2
+        H2 = torch.clamp(H2, eps, 1)
+        H5_pred = torch.cat([H2, s2], 1).detach() - self.denoise_2(torch.cat([H2, s2], 1))
+        H5_pred = torch.clamp(H5_pred, eps, 1)
+        H3 = H5_pred[:, :3, :, :]
+        return H2,H3
+

multi_read_data.py

@@ -0,0 +1,47 @@
+import numpy as np
+import torch
+import torch.utils.data
+from PIL import Image
+import torchvision.transforms as transforms
+import os
+
+class DataLoader(torch.utils.data.Dataset):
+    def __init__(self, img_dir, task):
+        self.low_img_dir = img_dir
+        self.task = task
+        self.train_low_data_names = []
+        self.train_target_data_names = []
+
+        for root, dirs, names in os.walk(self.low_img_dir):
+            for name in names:
+                self.train_low_data_names.append(os.path.join(root, name))
+
+        self.train_low_data_names.sort()
+        self.count = len(self.train_low_data_names)
+        transform_list = []
+        transform_list += [transforms.ToTensor()]
+        self.transform = transforms.Compose(transform_list)
+
+
+    def load_images_transform(self, file):
+
+        im = Image.open(file).convert('RGB')
+        img_norm = self.transform(im).numpy()
+        img_norm = np.transpose(img_norm, (1, 2, 0))
+        return img_norm
+
+
+    def __getitem__(self, index):
+
+        low = self.load_images_transform(self.train_low_data_names[index])
+        low = np.asarray(low, dtype=np.float32)
+        low = np.transpose(low[:, :, :], (2, 0, 1))
+        img_name = self.train_low_data_names[index].split('\\')[-1]
+
+
+
+
+        return torch.from_numpy(low),img_name
+
+    def __len__(self):
+        return self.count

test.py

@@ -0,0 +1,89 @@
+import os
+import sys
+import numpy as np
+import torch
+import argparse
+import logging
+import torch.utils
+from PIL import Image
+from torch.autograd import Variable
+from model import Finetunemodel
+from multi_read_data import DataLoader
+from thop import profile
+
+
+
+root_dir = os.path.abspath('../')
+sys.path.append(root_dir)
+
+parser = argparse.ArgumentParser("ZERO-IG")
+parser.add_argument('--data_path_test_low', type=str, default='./data',
+                    help='location of the data corpus')
+parser.add_argument('--save', type=str,
+                    default='./results/',
+                    help='location of the data corpus')
+parser.add_argument('--model_test', type=str,
+                    default='./model',
+                    help='location of the data corpus')
+parser.add_argument('--gpu', type=int, default=0, help='gpu device id')
+parser.add_argument('--seed', type=int, default=2, help='random seed')
+
+args = parser.parse_args()
+save_path = args.save
+os.makedirs(save_path, exist_ok=True)
+
+log_format = '%(asctime)s %(message)s'
+logging.basicConfig(stream=sys.stdout, level=logging.INFO,
+                    format=log_format, datefmt='%m/%d %I:%M:%S %p')
+mertic = logging.FileHandler(os.path.join(args.save, 'log.txt'))
+mertic.setFormatter(logging.Formatter(log_format))
+logging.getLogger().addHandler(mertic)
+
+logging.info("train file name = %s", os.path.split(__file__))
+TestDataset = DataLoader(img_dir=args.data_path_test_low,task='test')
+test_queue = torch.utils.data.DataLoader(TestDataset, batch_size=1, pin_memory=True, num_workers=0, shuffle=False)
+
+
+def save_images(tensor):
+    image_numpy = tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)))
+    im = np.clip(image_numpy * 255.0, 0, 255.0).astype('uint8')
+    return im
+
+def calculate_model_parameters(model):
+    return sum(p.numel() for p in model.parameters())
+
+def calculate_model_flops(model, input_tensor):
+    flops, _ = profile(model, inputs=(input_tensor,))
+    flops_in_gigaflops = flops / 1e9  # Convert FLOPs to gigaflops (G)
+    return flops_in_gigaflops
+
+def main():
+    if not torch.cuda.is_available():
+        print('no gpu device available')
+        sys.exit(1)
+
+    model = Finetunemodel(args.model_test)
+    model = model.cuda()
+    model.eval()
+    # Calculate model size
+    total_params = calculate_model_parameters(model)
+    print("Total number of parameters: ", total_params)
+    for p in model.parameters():
+        p.requires_grad = False
+    with torch.no_grad():
+        for _, (input,  img_name) in enumerate(test_queue):
+            input = Variable(input, volatile=True).cuda()
+            input_name = img_name[0].split('/')[-1].split('.')[0]
+            enhance,output = model(input)
+            input_name = '%s' % (input_name)
+            enhance=save_images(enhance)
+            output = save_images(output)
+            os.makedirs(args.save + '/result', exist_ok=True)
+            Image.fromarray(output).save(args.save + '/result/' +input_name + '_denoise' + '.png', 'PNG')
+            Image.fromarray(enhance).save(args.save + '/result/'+ input_name + '_enhance'  + '.png', 'PNG')
+    torch.set_grad_enabled(True)
+
+
+if __name__ == '__main__':
+    main()

train.py

@@ -0,0 +1,138 @@
+import os
+import sys
+import time
+import glob
+import numpy as np
+import utils
+from PIL import Image
+import logging
+import argparse
+import torch.utils
+import torch.backends.cudnn as cudnn
+from torch.autograd import Variable
+from model import *
+from multi_read_data import DataLoader
+
+
+parser = argparse.ArgumentParser("ZERO-IG")
+parser.add_argument('--batch_size', type=int, default=1, help='batch size')
+parser.add_argument('--cuda', default=True, type=bool, help='Use CUDA to train model')
+parser.add_argument('--gpu', type=str, default='0', help='gpu device id')
+parser.add_argument('--seed', type=int, default=2, help='random seed')
+parser.add_argument('--epochs', type=int, default=2001, help='epochs')
+parser.add_argument('--lr', type=float, default=0.0003, help='learning rate')
+parser.add_argument('--save', type=str, default='./EXP/', help='location of the data corpus')
+parser.add_argument('--model_pretrain', type=str,default='',help='location of the data corpus')
+
+args = parser.parse_args()
+
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+
+args.save = args.save + '/' + 'Train-{}'.format(time.strftime("%Y%m%d-%H%M%S"))
+utils.create_exp_dir(args.save, scripts_to_save=glob.glob('*.py'))
+model_path = args.save + '/model_epochs/'
+os.makedirs(model_path, exist_ok=True)
+image_path = args.save + '/image_epochs/'
+os.makedirs(image_path, exist_ok=True)
+
+log_format = '%(asctime)s %(message)s'
+logging.basicConfig(stream=sys.stdout, level=logging.INFO,
+                    format=log_format, datefmt='%m/%d %I:%M:%S %p')
+fh = logging.FileHandler(os.path.join(args.save, 'log.txt'))
+fh.setFormatter(logging.Formatter(log_format))
+logging.getLogger().addHandler(fh)
+
+logging.info("train file name = %s", os.path.split(__file__))
+
+if torch.cuda.is_available():
+    if args.cuda:
+        torch.set_default_tensor_type('torch.cuda.FloatTensor')
+    if not args.cuda:
+        print("WARNING: It looks like you have a CUDA device, but aren't " +
+              "using CUDA.\nRun with --cuda for optimal training speed.")
+        torch.set_default_tensor_type('torch.FloatTensor')
+else:
+    torch.set_default_tensor_type('torch.FloatTensor')
+
+
+def save_images(tensor):
+    image_numpy = tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)))
+    im = np.clip(image_numpy * 255.0, 0, 255.0).astype('uint8')
+    return im
+
+
+def main():
+    if not torch.cuda.is_available():
+        logging.info('no gpu device available')
+        sys.exit(1)
+
+    np.random.seed(args.seed)
+    cudnn.benchmark = True
+    torch.manual_seed(args.seed)
+    cudnn.enabled = True
+    torch.cuda.manual_seed(args.seed)
+    logging.info('gpu device = %s' % args.gpu)
+    logging.info("args = %s", args)
+
+
+
+    model =Network()
+    utils.save(model, os.path.join(args.save, 'initial_weights.pt'))
+    model.enhance.in_conv.apply(model.enhance_weights_init)
+    model.enhance.conv.apply(model.enhance_weights_init)
+    model.enhance.out_conv.apply(model.enhance_weights_init)
+    model = model.cuda()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), weight_decay=3e-4)
+    MB = utils.count_parameters_in_MB(model)
+    logging.info("model size = %f", MB)
+    print(MB)
+    train_low_data_names = './data/1'
+    TrainDataset = DataLoader(img_dir=train_low_data_names, task='train')
+
+    test_low_data_names = './data/1'
+    TestDataset = DataLoader(img_dir=test_low_data_names, task='test')
+
+    train_queue = torch.utils.data.DataLoader(
+        TrainDataset, batch_size=args.batch_size,
+        pin_memory=True, num_workers=0, shuffle=False, generator=torch.Generator(device='cuda'))
+    test_queue = torch.utils.data.DataLoader(
+        TestDataset, batch_size=1,
+        pin_memory=True, num_workers=0, shuffle=False, generator=torch.Generator(device='cuda'))
+
+    total_step = 0
+    model.train()
+    for epoch in range(args.epochs):
+        losses = []
+        for idx, (input, img_name) in enumerate(train_queue):
+            total_step += 1
+            input = Variable(input, requires_grad=False).cuda()
+            optimizer.zero_grad()
+            optimizer.param_groups[0]['capturable'] = True
+            loss = model._loss(input)
+            loss.backward()
+            nn.utils.clip_grad_norm_(model.parameters(), 5)
+            optimizer.step()
+            losses.append(loss.item())
+            logging.info('train-epoch %03d %03d %f', epoch, idx, loss)
+        logging.info('train-epoch %03d %f', epoch, np.average(losses))
+        utils.save(model, os.path.join(model_path, 'weights_%d.pt' % epoch))
+
+        if epoch % 50 == 0 and total_step != 0:
+            model.eval()
+            with torch.no_grad():
+                for idx, (input, img_name) in enumerate(test_queue):
+                    input = Variable(input, volatile=True).cuda()
+                    image_name = img_name[0].split('/')[-1].split('.')[0]
+                    L_pred1,L_pred2,L2,s2,s21,s22,H2,H11,H12,H13,s13,H14,s14,H3,s3,H3_pred,H4_pred,L_pred1_L_pred2_diff,H13_H14_diff,H2_blur,H3_blur= model(input)
+                    input_name = '%s' % (image_name)
+                    H3 = save_images(H3)
+                    H2= save_images(H2)
+                    os.makedirs(args.save + '/result/denoise/', exist_ok=True)
+                    os.makedirs(args.save + '/result/enhance/', exist_ok=True)
+                    Image.fromarray(H3).save(args.save + '/result/denoise/' + input_name+'_denoise_'+str(epoch)+'.png', 'PNG')
+                    Image.fromarray(H2).save(args.save + '/result/enhance/' +input_name+'_enhance_'+str(epoch)+'.png', 'PNG')
+
+
+if __name__ == '__main__':
+    main()

utils.py

@@ -0,0 +1,141 @@
+import os
+import numpy as np
+import torch
+import shutil
+from torch.autograd import Variable
+import matplotlib.pyplot as plt
+from PIL import Image
+
+
+
+def pair_downsampler(img):
+    # img has shape B C H W
+    c = img.shape[1]
+    filter1 = torch.FloatTensor([[[[0, 0.5], [0.5, 0]]]]).to(img.device)
+    filter1 = filter1.repeat(c, 1, 1, 1)
+    filter2 = torch.FloatTensor([[[[0.5, 0], [0, 0.5]]]]).to(img.device)
+    filter2 = filter2.repeat(c, 1, 1, 1)
+    output1 = torch.nn.functional.conv2d(img, filter1, stride=2, groups=c)
+    output2 = torch.nn.functional.conv2d(img, filter2, stride=2, groups=c)
+    return output1,output2
+
+def gauss_cdf(x):
+    return 0.5*(1+torch.erf(x/torch.sqrt(torch.tensor(2.))))
+
+def gauss_kernel(kernlen=21,nsig=3,channels=1):
+    interval=(2*nsig+1.)/(kernlen)
+    x=torch.linspace(-nsig-interval/2.,nsig+interval/2.,kernlen+1,).cuda()
+    #kern1d=torch.diff(torch.erf(x/math.sqrt(2.0)))/2.0
+    kern1d=torch.diff(gauss_cdf(x))
+    kernel_raw=torch.sqrt(torch.outer(kern1d,kern1d))
+    kernel=kernel_raw/torch.sum(kernel_raw)
+    #out_filter=kernel.unsqueeze(2).unsqueeze(3).repeat(1,1,channels,1)
+    out_filter=kernel.view(1,1,kernlen,kernlen)
+    out_filter = out_filter.repeat(channels,1,1,1)
+    return  out_filter
+
+class LocalMean(torch.nn.Module):
+    def __init__(self, patch_size=5):
+        super(LocalMean, self).__init__()
+        self.patch_size = patch_size
+        self.padding = self.patch_size // 2
+
+    def forward(self, image):
+        image = torch.nn.functional.pad(image, (self.padding, self.padding, self.padding, self.padding), mode='reflect')
+        patches = image.unfold(2, self.patch_size, 1).unfold(3, self.patch_size, 1)
+        return patches.mean(dim=(4, 5))
+
+def blur(x):
+    device = x.device
+    kernel_size = 21
+    padding = kernel_size // 2
+    kernel_var = gauss_kernel(kernel_size, 1, x.size(1)).to(device)
+    x_padded = torch.nn.functional.pad(x, (padding, padding, padding, padding), mode='reflect')
+    return torch.nn.functional .conv2d(x_padded, kernel_var, padding=0, groups=x.size(1))
+
+def padr_tensor(img):
+    pad=2
+    pad_mod=torch.nn.ConstantPad2d(pad,0)
+    img_pad=pad_mod(img)
+    return img_pad
+
+def calculate_local_variance(train_noisy):
+    b,c,w,h=train_noisy.shape
+    avg_pool = torch.nn.AvgPool2d(kernel_size=5,stride=1,padding=2)
+    noisy_avg= avg_pool(train_noisy)
+    noisy_avg_pad=padr_tensor(noisy_avg)
+    train_noisy=padr_tensor(train_noisy)
+    unfolded_noisy_avg=noisy_avg_pad.unfold(2,5,1).unfold(3,5,1)
+    unfolded_noisy=train_noisy.unfold(2,5,1).unfold(3,5,1)
+    unfolded_noisy_avg=unfolded_noisy_avg.reshape(unfolded_noisy_avg.shape[0],-1,5,5)
+    unfolded_noisy=unfolded_noisy.reshape(unfolded_noisy.shape[0],-1,5,5)
+    noisy_diff_squared=(unfolded_noisy-unfolded_noisy_avg)**2
+    noisy_var=torch.mean(noisy_diff_squared,dim=(2,3))
+    noisy_var=noisy_var.view(b,c,w,h)
+    return noisy_var
+
+def count_parameters_in_MB(model):
+    return np.sum(np.prod(v.size()) for name, v in model.named_parameters() if "auxiliary" not in name)/1e6
+
+
+
+def save_checkpoint(state, is_best, save):
+  filename = os.path.join(save, 'checkpoint.pth.tar')
+  torch.save(state, filename)
+  if is_best:
+    best_filename = os.path.join(save, 'model_best.pth.tar')
+    shutil.copyfile(filename, best_filename)
+
+
+def save(model, model_path):
+  torch.save(model.state_dict(), model_path)
+
+
+def load(model, model_path):
+  model.load_state_dict(torch.load(model_path))
+
+def drop_path(x, drop_prob):
+  if drop_prob > 0.:
+    keep_prob = 1.-drop_prob
+    mask = Variable(torch.cuda.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob))
+    x.div_(keep_prob)
+    x.mul_(mask)
+  return x
+
+def create_exp_dir(path, scripts_to_save=None):
+  if not os.path.exists(path):
+    os.makedirs(path,exist_ok=True)
+  print('Experiment dir : {}'.format(path))
+
+  if scripts_to_save is not None:
+    os.makedirs(os.path.join(path, 'scripts'),exist_ok=True)
+    for script in scripts_to_save:
+      dst_file = os.path.join(path, 'scripts', os.path.basename(script))
+      shutil.copyfile(script, dst_file)
+
+def show_pic(pic, name,path):
+    pic_num = len(pic)
+    for i in range(pic_num):
+        img = pic[i]
+        image_numpy = img[0].cpu().float().numpy()
+        if image_numpy.shape[0]==3:
+            image_numpy = (np.transpose(image_numpy, (1, 2, 0)))
+            im = Image.fromarray(np.clip(image_numpy * 255.0, 0, 255.0).astype('uint8'))
+            img_name = name[i]
+            plt.subplot(5, 6, i + 1)
+            plt.xlabel(str(img_name))
+            plt.xticks([])
+            plt.yticks([])
+            plt.imshow(im)
+        elif image_numpy.shape[0]==1:
+            im = Image.fromarray(np.clip(image_numpy[0] * 255.0, 0, 255.0).astype('uint8'))
+            img_name = name[i]
+            plt.subplot(5, 6, i + 1)
+            plt.xlabel(str(img_name))
+            plt.xticks([])
+            plt.yticks([])
+            plt.imshow(im,plt.cm.gray)
+    plt.savefig(path)
+
+
+