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

第一步：HEP介绍

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

第二步：HEP网络结构

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

第三步：模型代码展示

python 复制代码

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:
    pass

parser = 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)
        # Refelectance
        x0 = 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)
        # Outputs
        R = 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))

第四步：运行

第五步：整个工程的内容

代码的下载路径 （新窗口打开链接）：基于深度学习神经网络的AI弱光图像增强HEP系统源码

有问题可以私信或者留言，有问必答