基于深度学习神经网络的AI弱光图像增强HEP系统源码

第一步：HEP介绍

基于深度学习的微光图像增强方法通常需要大量的成对训练数据，这在现实世界中是不切实际的。最近，已经探索了无监督的方法来消除对成对训练数据的依赖。然而，由于缺乏先验，它们在不同的现实世界场景中表现不稳定。为了解决这个问题，文章提出了一种基于有效先验的无监督微光图像增强方法，称为直方图均衡先验（HEP）。文章的工作是受到了一个有趣的观察结果的启发，即直方图均衡增强图像的特征图和地面实况是相似的。具体而言，文章制定HEP以提供丰富的纹理和亮度信息。它嵌入到发光模块（LUM）中，有助于将低光图像分解为照度图和反射率图，反射率图可以被视为恢复的图像。然而，基于Retinex理论的推导表明，反射率图受到了噪声的污染。文章引入了一种噪声去纠缠模块（NDM），在未配对干净图像的可靠帮助下，去纠缠反射图中的噪声和内容。在直方图均衡先验和噪声解纠缠的指导下，我们的方法可以恢复更精细的细节，并且能够在现实世界的低光场景中抑制噪声。大量的实验表明，我们的方法与最先进的无监督微光增强算法相比表现良好，甚至与最新的有监督算法相匹配。

第二步：HEP网络结构

算法整体流程如下图，结合提亮和降噪。

第三步：模型代码展示

from abc import ABC
from thop import profile
from torch import nn
import torch
import argparse
from utils import get_config
from models.NDM_model import Conv2dBlock
try:from itertools import izip as zip
except ImportError:passparser = argparse.ArgumentParser()
parser.add_argument('--denoise_config', type=str, default='./configs/unit_NDM.yaml', help="denoise net configuration")
parser.add_argument('--light_config', type=str, default='configs/unit_LUM.yaml', help='Path to the config file.')
parser.add_argument('--input_folder', type=str, default='./test_images', help="input image path")
parser.add_argument('--output_folder', type=str, default='./NDM_results', help="output image path")
parser.add_argument('--denoise_checkpoint', type=str, default='./checkpoints/NDM_LOL.pt', help="checkpoint of denoise")
parser.add_argument('--light_checkpoint', type=str, default='./checkpoints/LUM_LOL.pth', help="checkpoint of light")
opts = parser.parse_args()class DecomNet(nn.Module, ABC):def __init__(self, params):super(DecomNet, self).__init__()self.norm = params['norm']self.activ = params['activ']self.pad_type = params['pad_type']#self.conv0 = Conv2dBlock(4, 32, 3, 1, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)self.conv1 = Conv2dBlock(4, 64, 9, 1, 4, norm=self.norm, activation='none', pad_type=self.pad_type)self.conv2 = Conv2dBlock(64, 64, 3, 1, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)self.conv3 = Conv2dBlock(64, 128, 3, 2, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)self.conv4 = Conv2dBlock(128, 128, 3, 1, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)self.conv5 = nn.ConvTranspose2d(128, 64, 3, 2, 1, 1)self.activation = nn.ReLU(inplace=True)self.conv6 = Conv2dBlock(128, 64, 3, 1, 1, norm=self.norm, activation=self.activ, pad_type=self.pad_type)self.conv7 = Conv2dBlock(96, 64, 3, 1, 1, norm=self.norm, activation='none', pad_type=self.pad_type)self.conv8 = Conv2dBlock(64, 4, 3, 1, 1, norm=self.norm, activation='none', pad_type=self.pad_type)def forward(self, input_im):input_max = torch.max(input_im, dim=1, keepdim=True)[0]image = torch.cat((input_max, input_im), dim=1)# Refelectancex0 = self.conv0(image)# print('x0:', x0.shape)x1 = self.conv1(image)# print('x1:', x1.shape)x2 = self.conv2(x1)# print('x2:', x2.shape)x3 = self.conv3(x2)# print('x3:', x3.shape)x4 = self.conv4(x3)# print('x4:', x4.shape)x5 = self.conv5(x4)x5 = self.activation(x5)# print('x5:', x5.shape)cat5 = torch.cat((x5, x2), dim=1)x6 = self.conv6(cat5)# print('x6:', x6.shape)cat6 = torch.cat((x6, x0), dim=1)x7 = self.conv7(cat6)# print('x7:', x7.shape)x8 = self.conv8(x7)# print('x8:', x8.shape)# OutputsR = torch.sigmoid(x8[:, 0:3, :, :])L = torch.sigmoid(x8[:, 3:4, :, :])return R, L
if __name__ == "__main__":light_config = get_config(opts.light_config)model = DecomNet(light_config)input = torch.randn(1, 3, 256, 256)flops, params = profile(model, inputs=(input, ))print("flops:{}".format(flops))print("params:{}".format(params))

第四步：运行