基于深度学习神经网络的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系统源码

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

相关推荐
冬奇Lab17 小时前
Workflow 系列(03):状态管理——持久化、幂等性与版本绑定
人工智能·工作流引擎
冬奇Lab17 小时前
每日一个开源项目(第146篇):openpilot - 开源自动驾驶辅助系统,曾在 Consumer Reports 评测中超过特斯拉 Autopilot
人工智能·开源·自动驾驶
吴佳浩18 小时前
AI 工程师知识地图:模型格式、框架、部署工具一次讲明白
人工智能·aigc·ai编程
IT_陈寒19 小时前
Java的Date类又坑了我一次,改用时间戳真香
前端·人工智能·后端
码农胖大海19 小时前
AI额度不够用的解决方案
人工智能
后端小肥肠19 小时前
小红书虚拟商品怎么做?我先用 Skill 跑通了壁纸品类
人工智能·aigc·agent
feiyu_gao19 小时前
从零搭建个人 AI 工作台:一个管理者的 3 个月实验
人工智能·aigc·团队管理
Lihua奏20 小时前
从单核到多核:CPU为什么不能再只靠提频变快
深度学习
程序员cxuan20 小时前
一句话,让你用上 GPT-5.6
人工智能·后端·程序员