低照度图像增强模型——损失函数整理汇总

Vaeeeeeee

于 2024-04-12 20:08:32 发布

阅读量1.1k

点赞数 21

文章标签：损失函数低照度模型 Pytorch

本文链接：https://blog.csdn.net/m0_46366547/article/details/137337592

版权

文章目录

前言
符号说明
1、L1损失函数 / 平均绝对误差（MAE）
2、SSIM损失函数
3、平滑度损失（Smoothness-Loss）⭐️
3、颜色损失（Color Constancy Loss）
4、感知损失（perceptual loss）
5、曝光损失（Exposure Loss）
总结

前言

损失函数在低照度增强模型设计中扮演着至关重要的角色。它是评估模型性能的关键指标，通过衡量模型输出与真实图像之间的差异，指导模型在训练过程中不断优化。在低照度图像增强任务中，损失函数的选择和设计直接影响到模型对图像细节的恢复能力和整体质量的提升。因此，在模型设计过程中，需要根据具体任务需求和图像特点选择合适的损失函数，以确保模型能够有效地实现低照度图像的增强。
本文主要针对近几年的低照度模型中常使用的损失函数做个简单的整理。如果需要构造新的、自己的损失函数，方便有个参考。📚📚📚

符号说明

$Ground\quad Truth,（GT图像）：曝光良好的图像$
$\quad img：模型输出的图像$
$I\_img:光照分量图像$
$R\_img:反射分量图像$

1、L1损失函数 / 平均绝对误差（MAE）

$非常常见的损失函数，在绝大部分的深度学习模型当中都要用到。在低照度增强模型中一般用来衡量经过模型增\\ 强后的输出图像与GT图像之间的误差。$
$\begin{aligned} L1_{loss}=|out\_img-gt\_img| \end{aligned}$
$用法：$

# 第一种：
import torch.nn.functional as F
F.l1_loss(out_image ， gt_image)


# 第二种：
import torch.nn as nn
nn.L1Loss()(out_image ， gt_image)


# 第三种：
def MAELoss(out_image, gt_image):
    return torch.mean(torch.abs(out_image - gt_image))

2、SSIM损失函数

$同样属于经常使用的损失函数，一般衡量GT图像和输出图像之间的差距。SSIM是综合亮度、对比度和结构\\3个方面来评估图像质量的，而SSIM损失去掉了亮度这一维度，模型训练过程中可以有助于很好保留结构和\\纹理。$
$\begin{aligned} L_{ms-ssim}=1-\prod_{m=1}^{M}(\frac{2\mu_{p}\mu_{g}+c_{1}}{\mu_{g}^{2}+\mu_{g}^{2}+c_{1}})^{\beta_{m}}(\frac{2\sigma_{pg}+c_{2}}{\sigma_{p}^{2}+\sigma_{g}^{2}+c_{2}})^{\gamma_{m}}\\ M:表示不同尺度;\mu_{p},\mu_{g}分别表示out\_img和GT的均值；\\ \sigma_{p},\sigma_{g}分别表示out\_img和GT的标准差；\sigma_{pg}表示二者之间的协方差\\ \beta_{m}和\gamma_{m}，c_{1}和c_{2}都为常数 \end{aligned}$
$代码实现：$

from pytorch_msssim import ms_ssim
def MS_SSIMLoss(out_image, gt_image):
    return 1 - ms_ssim(out_image, gt_image, data_range=1, size_average=True)

3、平滑度损失（Smoothness-Loss）⭐️

$顾名思义就是让图片像德芙一样的丝滑，并且保持相邻像素之间的关系$ ~（😃） $因为通常认为自然图像的内容变化就是平缓的，原理实现上就是\textbf{尽量最小化像素与邻域的梯度}。\\ 所以也被称为TV\_loss。$
$out\_img上，更多的用来约束\textbf{估计的光照量}，因为有个先验知识就是\\ 光照的变化是连续的，渐变的，一个理想的光照分量在纹理细节上尽可能的平滑，呈现出“局部光滑”的特点。$
$对于彩色图像，可以在颜色通道上分别计算总变分，然后最小化各通道总和。$

$\begin{aligned} L_{TV}=\frac1N\Sigma_{n=1}^N\Sigma_{c\in\xi}(|\nabla_xout\_img_n^c|+|\nabla_yout\_img_n^c|)^2,\xi=\{R,G,B\} \end{aligned}\\N是迭代次数；\nabla_x和\nabla_y分别表示水平和垂直方向上的梯度值$

$代码：$

class TVLoss(nn.Module):
    def __init__(self):
        super(TVLoss,self).__init__()

    def forward(self,x,weight_map=None):
        self.h_x = x.size()[2]
        self.w_x = x.size()[3]
        self.batch_size = x.size()[0]
        if weight_map is None:
            self.TVLoss_weight=(1, 1)
        else:
            self.TVLoss_weight = self.compute_weight(weight_map)

        count_h = self._tensor_size(x[:,:,1:,:])
        count_w = self._tensor_size(x[:,:,:,1:])

        h_tv = (self.TVLoss_weight[0]*torch.abs((x[:,:,1:,:]-x[:,:,:self.h_x-1,:]))).sum()
        w_tv = (self.TVLoss_weight[1]*torch.abs((x[:,:,:,1:]-x[:,:,:,:self.w_x-1]))).sum()
        # print(self.TVLoss_weight[0],self.TVLoss_weight[1])
        return (h_tv/count_h+w_tv/count_w)/self.batch_size

    def _tensor_size(self,t):
        return t.size()[1]*t.size()[2]*t.size()[3]

    def compute_weight(self, img):
        gradx = torch.abs(img[:, :, 1:, :] - img[:, :, :self.h_x-1, :])
        grady = torch.abs(img[:, :, :, 1:] - img[:, :, :, :self.w_x-1])
        TVLoss_weight_x = torch.div(1,torch.exp(gradx))
        TVLoss_weight_y = torch.div(1, torch.exp(grady))

        return TVLoss_weight_x, TVLoss_weight_y

# 简洁版
class L_TV(nn.Module):
    def __init__(self):
        super(L_TV,self).__init__()

    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 2*(h_tv/count_h+w_tv/count_w)/batch_size


# DSLR中的
def tv_loss(img):
    """
    Compute total variation loss.
    Inputs:
    - img: PyTorch Variable of shape (1, 3, H, W) holding an input image.
    - tv_weight: Scalar giving the weight w_t to use for the TV loss.
    Returns:
    - loss: PyTorch Variable holding a scalar giving the total variation loss
      for img weighted by tv_weight.
    """
    b,c,h,w_ = img.size()
    w_variance = torch.sum(torch.pow(img[:,:,:,:-1] - img[:,:,:,1:], 2))/b
    h_variance = torch.sum(torch.pow(img[:,:,:-1,:] - img[:,:,1:,:], 2))/b
    loss = (h_variance + w_variance) / 2
    return loss

$求梯度来给光照分量I的梯度图\\ 分配权重，使得反射分量 R 上较为平滑的区域对应到光照分量I上同样也应该尽可能平滑。$
$\begin{aligned} \mathcal{L}_{is}=\sum_{i=low,normal}||\nabla I_i\circ\exp(-\lambda_g\nabla R_i)||, \end{aligned}\\$
$代码$

# RetinexNet
import torch
import torch.nn.functional as F
class Smoothloss():
	def __init__(self):
        super(Smoothloss, self).__init__()
        
	def gradient(self, input_tensor, direction):
	    self.smooth_kernel_x = torch.FloatTensor([[0, 0], [-1, 1]]).view((1, 1, 2, 2)).cuda()
	    self.smooth_kernel_y = torch.transpose(self.smooth_kernel_x, 2, 3)
	
	    if direction == "x":
	        kernel = self.smooth_kernel_x
	    elif direction == "y":
	        kernel = self.smooth_kernel_y
	    grad_out = torch.abs(F.conv2d(input_tensor, kernel,
	                                  stride=1, padding=1))
	    return grad_out
	    
	def ave_gradient(self, input_tensor, direction):
	    return F.avg_pool2d(self.gradient(input_tensor, direction),
	                        kernel_size=3, stride=1, padding=1)
	
	def smooth(self, I_img, R_img):
	    R_img = 0.299*R_img[:, 0, :, :] + 0.587*R_img[:, 1, :, :] + 0.114*R_img[:, 2, :, :]  # 转换到YUV空间中的Y通道
	    R_img = torch.unsqueeze(R_img, dim=1)
	    return torch.mean(self.gradient(I_img, "x") * torch.exp(-10 * self.ave_gradient(R_img, "x")) +
	                      self.gradient(I_img, "y") * torch.exp(-10 * self.ave_gradient(R_img, "y")))

3、颜色损失（Color Constancy Loss）

$用于纠正$ $out\_img$ $中的潜在色偏$ $，用于保证和$ $GT\_img$ $之间色彩一致，尽可能避免色调失真的问题。$
$Z ero - D CE$ $中的$ $\mathcal{L}_{color}$ 为：
$\mathcal{L}_{color}=\sum_{\forall(p,q)\in\varepsilon}(J^p-J^q)^2,\varepsilon=\{(R,G),(R,B),(G,B)\}\\ J^{p}表示增强图像中p通道的平均值\\ （p,q）表示一对颜色通道$
$代码（适合无监督训练）：$

class L_color(nn.Module):
 
    def __init__(self):
        super(L_color, self).__init__()
 
    def forward(self, x ):
 
        b,c,h,w = x.shape
 
        mean_rgb = torch.mean(x,[2,3],keepdim=True)
        mr,mg, mb = torch.split(mean_rgb, 1, dim=1)
        Drg = torch.pow(mr-mg,2)
        Drb = torch.pow(mr-mb,2)
        Dgb = torch.pow(mb-mg,2)
        k = torch.pow(torch.pow(Drg,2) + torch.pow(Drb,2) + torch.pow(Dgb,2),0.5)
        return k

$D I V F u s i o n$ $中的$ $\mathcal{L}_{color}$ $为：$
$\mathcal{L}_{color}=angle(I_{f}^{i},I_{vi}^{i}),\quad i\in\{R,G,B\}\\ angel表示余弦相似度$
$代码：（适合有监督训练）$

def angle(a, b):
    vector = torch.mul(a, b)
    up     = torch.sum(vector)
    down   = torch.sqrt(torch.sum(torch.square(a))) * torch.sqrt(torch.sum(torch.square(b)))
    theta  = torch.acos(up/down) # 弧度制
    return theta
def color_loss(out_image, gt_image): # 颜色损失  希望增强前后图片的颜色一致性 (b,c,h,w)
    loss = torch.mean(angle(out_image[:,0,:,:],gt_image[:,0,:,:]) + 
                      angle(out_image[:,1,:,:],gt_image[:,1,:,:]) +
                      angle(out_image[:,2,:,:],gt_image[:,2,:,:]))
    return loss

4、感知损失（perceptual loss）

$一种基于神经网络特征的损失函数，可以提高结果的视觉质量。它被定义为增强的结果的特征表示与\\相应的 GT\_img之间的欧氏距离，即通过比较目标图像和增强图像在\textbf{高层特征图上的差异}\\来度量图像的相似性，特征表示通常是从 ImageNet 数据集上预训练的 VGG网络中提取的。$
$\mathcal{L}_{perceptual}=\frac{1}{C_jH_jW_j}\|\phi_j(gt\_img)-\phi_j(out\_img)\|_2^2\\ C_{j}H_{j}W_{j};表示第j层的特征图的大小\\ \phi表示损失网络，一般使用VGG$
$代码：$

import torch
import torch.nn.functional as F
from torchvision import models
# 加载预训练的VGG模型
vgg = models.vgg19(pretrained=True).features
# 将模型设置为评估模式
vgg.eval()
def perceptual_loss(out_img, gt_img):
    # 计算输入和目标图像在VGG特征图上的差异
    input_features = vgg(out_img)
    target_features = vgg(gt_img)
    loss = F.mse_loss(input_features, target_features)
    return loss


# 如果要从指定路径加载模型权重
def load_vgg(weights_path : str):
    vgg = models.vgg16(pretrained=False)         # 加载 VGG16 / VGG19 模型的特征提取部分，不使用预训练权重
    state_dict = torch.load(weights_path)  	 			 # 加载权重
    vgg.load_state_dict(state_dict)   					 # 将加载的权重加载到模型中
    vgg = vgg.features
    return vgg

vgg = load_vgg(r'E:\预训练权重\vgg16-397923af.pth')
vgg.eval()			# 将模型设置为评估模式

def perceptual_loss(out_img, gt_img):
    # 计算输入和目标图像在 VGG 特征图上的差异
    input_features = vgg(out_img)
    target_features = vgg(gt_img)
    loss = F.mse_loss(input_features, target_features)
    return loss
perceptual_loss(a,b)

5、曝光损失（Exposure Loss）

$[0,1]区间的\\图片的平均灰度值一般在 0.6左右。所以可以通过计算增强图像与 GT图像之间的灰度均值\\的差值构造损失函数，一般对图像分块计算然后求总误差，可以\textbf{抑制曝光较差的区域，使其到达较好的曝光水平}。$
$L_{exposure}=\frac{1}{M}\sum_{k=1}^{M}|out\_img_{k}-E|\\ M表示分块区域的个数，分块大小一般取16×16\\ out\_img_{k}表示在第k块区域上的平均强度值$
$代码：$

class L_exp(nn.Module):
    def __init__(self,patch_size = 16,mean_val = 0.6):   # 如果图像像素值在0~255，mean_valy应该是255*0.6
        super(L_exp, self).__init__()
        # print(1)
        self.pool = nn.AvgPool2d(patch_size)
        self.mean_val = mean_val
        
    def forward(self, x ):
        b,c,h,w = x.shape
        x = torch.mean(x,1,keepdim=True)
        mean = self.pool(x)
 
        loss = torch.mean(torch.pow(mean- torch.FloatTensor([self.mean_val] ).cuda(),2))
        return loss

总结

各类低照度模型中几乎都用到了平滑损失函数作为整个损失函数的一部分，不仅可以对增强结果进行平滑约束，也可以对中间量光照量进行平滑约束，可以作为重点参考的损失函数。

Vaeeeeeee

关注

21
点赞
踩
34

收藏

觉得还不错? 一键收藏
0
评论
低照度图像增强模型——损失函数整理汇总

损失函数在低照度增强模型设计中扮演着至关重要的角色。它是评估模型性能的关键指标，通过衡量模型输出与真实图像之间的差异，指导模型在训练过程中不断优化。在低照度图像增强任务中，损失函数的选择和设计直接影响到模型对图像细节的恢复能力和整体质量的提升。本文主要针对近几年的低照度模型中常使用的损失函数做个简单的整理。如果需要构造新的、自己的损失函数，方便有个参考。
复制链接

扫一扫