colorlosspytorch实现

1.实现原理

最近在做图像增强相关的工作&＃xff0c;偶然间看到了这篇文章&＃xff0c;作者提出了一个损失叫做color-loss&＃xff0c;根据文章描述该方法是通过模糊输入图像与ground-Truth的纹理、内容&＃xff0c;仅仅保存图像的颜色信息实现图像颜色的校正。实现过程比较简单&＃xff0c;首先构建一个高斯模糊核&＃xff0c;然后利用高斯模糊核作为卷积核对图像进行卷积运算&＃xff0c;得到模糊后的图像&＃xff1b;然后计算输入图像与ground-Truth的MSE作为损失函数。
作者的github中有该模型的代码&＃xff0c;但是是用TensorFlow实现的。因为我的代码pytorch的&＃xff0c;所以自己重新改写了一下。在作者的代码中用到了深度可分离卷积&＃xff0c;在pytorch中我没有对其进行深度可分离操作。
算是个深度学习的小白吧&＃xff0c;有问题可以给我留言呀~~~

2.代码

// An highlighted block
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from math import exp, pi
import numpy as np
import cv2 as cv
import scipy.stats as st
import matplotlib.pyplot as pltdef gauss_kernel(kernlen&＃61;21, nsig&＃61;3, channels&＃61;1):interval &＃61; (2*nsig&＃43;1.)/(kernlen)x &＃61; np.linspace(-nsig-interval/2., nsig&＃43;interval/2., kernlen&＃43;1)kern1d &＃61; np.diff(st.norm.cdf(x))kernel_raw &＃61; np.sqrt(np.outer(kern1d, kern1d))kernel &＃61; kernel_raw/kernel_raw.sum()out_filter &＃61; np.array(kernel, dtype &＃61; np.float32)out_filter &＃61; out_filter.reshape((kernlen, kernlen))# out_filter &＃61; np.repeat(out_filter, channels, axis &＃61; 0)return out_filter # kernel_size&＃61;21class SeparableConv2d(nn.Module):def __init__(self):super(SeparableConv2d, self).__init__()kernel &＃61; gauss_kernel(21, 3, 3)kernel &＃61; torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0)## kernel_point &＃61; [[1.0]]## kernel_point &＃61; torch.FloatTensor(kernel_point).unsqueeze(0).unsqueeze(0)# kernel &＃61; torch.FloatTensor(kernel).expand(3, 3, 21, 21) # torch.expand(&＃xff09;向输入的维度前面进行扩充&＃xff0c;输入为三通道时&＃xff0c;将weight扩展为[3,3,21,21]## kernel_point &＃61; torch.FloatTensor(kernel_point).expand(3,3,1,1)self.weight &＃61; nn.Parameter(data&＃61;kernel, requires_grad&＃61;False)# self.pointwise &＃61; nn.Conv2d(1, 1, 1, 1, 0, 1, 1,bias&＃61;False) # 单通道时in_channels&＃61;1&＃xff0c;out_channels&＃61;1,三通道时&＃xff0c;in_channels&＃61;3, out_channels&＃61;3 卷积核为随机的## self.weight_point &＃61; nn.Parameter(data&＃61;kernel_point, requires_grad&＃61;False)def forward(self, img1):x &＃61; F.conv2d(img1, self.weight, groups&＃61;1,padding&＃61;10)## x &＃61; F.conv2d(x, self.weight_point, groups&＃61;1, padding&＃61;0) #卷积核为[1]# x &＃61; self.pointwise(x)return x
# plt.imshow(out_kernel)
# plt.imshow(out_kernel)