计算像素周围8个邻近点的损失

计算像素周围8个邻近点的损失

其中，Y和I分别为增强图像和弱光图像中a×a局部区域的平均像素值。A是局部区域的一侧，我们根据消融研究将其设置为4。ϕ(i)是四个相邻的邻居（从上、下、左、右），而ψ(i)是四个不相邻的邻居（从左上、右上、左下和右下）。α为0.5，因为非相邻邻居的权值不那么重要。

class L_spa8(nn.Module):
    def __init__(self, patch_size):
        super(L_spa8, self).__init__()
        # print(1)kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0)
        # Build conv kernels
        kernel_left = torch.FloatTensor( [[0,0,0],[-1,1,0],[0,0,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_right = torch.FloatTensor( [[0,0,0],[0,1,-1],[0,0,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_up = torch.FloatTensor( [[0,-1,0],[0,1, 0 ],[0,0,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_down = torch.FloatTensor( [[0,0,0],[0,1, 0],[0,-1,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_upleft = torch.FloatTensor( [[-1,0,0],[0,1,0],[0,0,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_upright = torch.FloatTensor( [[0,0,-1],[0,1,0],[0,0,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_loleft = torch.FloatTensor( [[0,0,0],[0,1,0],[-1,0,0]]).to(device).unsqueeze(0).unsqueeze(0)
        kernel_loright = torch.FloatTensor( [[0,0,0],[0,1,0],[0,0,-1]]).to(device).unsqueeze(0).unsqueeze(0)

        # convert to parameters
        self.weight_left = nn.Parameter(data=kernel_left, requires_grad=False)
        self.weight_right = nn.Parameter(data=kernel_right, requires_grad=False)
        self.weight_up = nn.Parameter(data=kernel_up, requires_grad=False)
        self.weight_down = nn.Parameter(data=kernel_down, requires_grad=False)
        self.weight_upleft = nn.Parameter(data=kernel_upleft, requires_grad=False)
        self.weight_upright = nn.Parameter(data=kernel_upright, requires_grad=False)
        self.weight_loleft = nn.Parameter(data=kernel_loleft, requires_grad=False)
        self.weight_loright = nn.Parameter(data=kernel_loright, requires_grad=False)

        # pooling layer
        self.pool = nn.AvgPool2d(patch_size) # default is 4

    def forward(self, org , enhance ):
        #b,c,h,w = org.shape

        org_mean = torch.mean(org,1,keepdim=True)
        enhance_mean = torch.mean(enhance,1,keepdim=True)

        org_pool =  self.pool(org_mean)
        enhance_pool = self.pool(enhance_mean)

        #weight_diff =torch.max(torch.FloatTensor([1]).to(device) + 10000*torch.min(org_pool - torch.FloatTensor([0.3]).to(device),torch.FloatTensor([0]).to(device)),torch.FloatTensor([0.5]).to(device))
        #E_1 = torch.mul(torch.sign(enhance_pool - torch.FloatTensor([0.5]).to(device)) ,enhance_pool-org_pool)


        # Original output
        D_org_letf = F.conv2d(org_pool , self.weight_left, padding=1)
        D_org_right = F.conv2d(org_pool , self.weight_right, padding=1)
        D_org_up = F.conv2d(org_pool , self.weight_up, padding=1)
        D_org_down = F.conv2d(org_pool , self.weight_down, padding=1)
        D_org_upleft = F.conv2d(org_pool , self.weight_upleft , padding=1)
        D_org_upright = F.conv2d(org_pool , self.weight_upright, padding=1)
        D_org_loleft = F.conv2d(org_pool , self.weight_loleft, padding=1)
        D_org_loright = F.conv2d(org_pool , self.weight_loright, padding=1)


        # Enhanced output
        D_enhance_letf = F.conv2d(enhance_pool , self.weight_left, padding=1)
        D_enhance_right = F.conv2d(enhance_pool , self.weight_right, padding=1)
        D_enhance_up = F.conv2d(enhance_pool , self.weight_up, padding=1)
        D_enhance_down = F.conv2d(enhance_pool , self.weight_down, padding=1)
        D_enhance_upleft = F.conv2d(enhance_pool, self.weight_upleft, padding=1)
        D_enhance_upright = F.conv2d(enhance_pool, self.weight_upright, padding=1)
        D_enhance_loleft = F.conv2d(enhance_pool, self.weight_loleft, padding=1)
        D_enhance_loright = F.conv2d(enhance_pool, self.weight_loright, padding=1)

        # Difference
        D_left = torch.pow(D_org_letf - D_enhance_letf,2)
        D_right = torch.pow(D_org_right - D_enhance_right,2)
        D_up = torch.pow(D_org_up - D_enhance_up,2)
        D_down = torch.pow(D_org_down - D_enhance_down,2)
        D_upleft = torch.pow(D_org_upleft - D_enhance_upleft,2)
        D_upright = torch.pow(D_org_upright - D_enhance_upright,2)
        D_loleft = torch.pow(D_org_loleft - D_enhance_loleft,2)
        D_loright = torch.pow(D_org_loright - D_enhance_loright,2)

        # Total difference
        E = (D_left + D_right + D_up +D_down) + 0.5 * (D_upleft + D_upright + D_loleft + D_loright)

        # E = 25*(D_left + D_right + D_up +D_down)

        return E