295 lines
8.2 KiB
Python
295 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
#############################################################
|
|
# File: esa.py
|
|
# Created Date: Tuesday April 28th 2022
|
|
# Author: Chen Xuanhong
|
|
# Email: chenxuanhongzju@outlook.com
|
|
# Last Modified: Thursday, 20th April 2023 9:28:06 am
|
|
# Modified By: Chen Xuanhong
|
|
# Copyright (c) 2020 Shanghai Jiao Tong University
|
|
#############################################################
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .layernorm import LayerNorm2d
|
|
|
|
|
|
def moment(x, dim=(2, 3), k=2):
|
|
assert len(x.size()) == 4
|
|
mean = torch.mean(x, dim=dim).unsqueeze(-1).unsqueeze(-1)
|
|
mk = (1 / (x.size(2) * x.size(3))) * torch.sum(torch.pow(x - mean, k), dim=dim)
|
|
return mk
|
|
|
|
|
|
class ESA(nn.Module):
|
|
"""
|
|
Modification of Enhanced Spatial Attention (ESA), which is proposed by
|
|
`Residual Feature Aggregation Network for Image Super-Resolution`
|
|
Note: `conv_max` and `conv3_` are NOT used here, so the corresponding codes
|
|
are deleted.
|
|
"""
|
|
|
|
def __init__(self, esa_channels, n_feats, conv=nn.Conv2d):
|
|
super(ESA, self).__init__()
|
|
f = esa_channels
|
|
self.conv1 = conv(n_feats, f, kernel_size=1)
|
|
self.conv_f = conv(f, f, kernel_size=1)
|
|
self.conv2 = conv(f, f, kernel_size=3, stride=2, padding=0)
|
|
self.conv3 = conv(f, f, kernel_size=3, padding=1)
|
|
self.conv4 = conv(f, n_feats, kernel_size=1)
|
|
self.sigmoid = nn.Sigmoid()
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, x):
|
|
c1_ = self.conv1(x)
|
|
c1 = self.conv2(c1_)
|
|
v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
|
|
c3 = self.conv3(v_max)
|
|
c3 = F.interpolate(
|
|
c3, (x.size(2), x.size(3)), mode="bilinear", align_corners=False
|
|
)
|
|
cf = self.conv_f(c1_)
|
|
c4 = self.conv4(c3 + cf)
|
|
m = self.sigmoid(c4)
|
|
return x * m
|
|
|
|
|
|
class LK_ESA(nn.Module):
|
|
def __init__(
|
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
|
):
|
|
super(LK_ESA, self).__init__()
|
|
f = esa_channels
|
|
self.conv1 = conv(n_feats, f, kernel_size=1)
|
|
self.conv_f = conv(f, f, kernel_size=1)
|
|
|
|
kernel_size = 17
|
|
kernel_expand = kernel_expand
|
|
padding = kernel_size // 2
|
|
|
|
self.vec_conv = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(1, kernel_size),
|
|
padding=(0, padding),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
self.vec_conv3x1 = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(1, 3),
|
|
padding=(0, 1),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
|
|
self.hor_conv = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(kernel_size, 1),
|
|
padding=(padding, 0),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
self.hor_conv1x3 = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(3, 1),
|
|
padding=(1, 0),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
|
|
self.conv4 = conv(f, n_feats, kernel_size=1)
|
|
self.sigmoid = nn.Sigmoid()
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, x):
|
|
c1_ = self.conv1(x)
|
|
|
|
res = self.vec_conv(c1_) + self.vec_conv3x1(c1_)
|
|
res = self.hor_conv(res) + self.hor_conv1x3(res)
|
|
|
|
cf = self.conv_f(c1_)
|
|
c4 = self.conv4(res + cf)
|
|
m = self.sigmoid(c4)
|
|
return x * m
|
|
|
|
|
|
class LK_ESA_LN(nn.Module):
|
|
def __init__(
|
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
|
):
|
|
super(LK_ESA_LN, self).__init__()
|
|
f = esa_channels
|
|
self.conv1 = conv(n_feats, f, kernel_size=1)
|
|
self.conv_f = conv(f, f, kernel_size=1)
|
|
|
|
kernel_size = 17
|
|
kernel_expand = kernel_expand
|
|
padding = kernel_size // 2
|
|
|
|
self.norm = LayerNorm2d(n_feats)
|
|
|
|
self.vec_conv = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(1, kernel_size),
|
|
padding=(0, padding),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
self.vec_conv3x1 = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(1, 3),
|
|
padding=(0, 1),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
|
|
self.hor_conv = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(kernel_size, 1),
|
|
padding=(padding, 0),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
self.hor_conv1x3 = nn.Conv2d(
|
|
in_channels=f * kernel_expand,
|
|
out_channels=f * kernel_expand,
|
|
kernel_size=(3, 1),
|
|
padding=(1, 0),
|
|
groups=2,
|
|
bias=bias,
|
|
)
|
|
|
|
self.conv4 = conv(f, n_feats, kernel_size=1)
|
|
self.sigmoid = nn.Sigmoid()
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, x):
|
|
c1_ = self.norm(x)
|
|
c1_ = self.conv1(c1_)
|
|
|
|
res = self.vec_conv(c1_) + self.vec_conv3x1(c1_)
|
|
res = self.hor_conv(res) + self.hor_conv1x3(res)
|
|
|
|
cf = self.conv_f(c1_)
|
|
c4 = self.conv4(res + cf)
|
|
m = self.sigmoid(c4)
|
|
return x * m
|
|
|
|
|
|
class AdaGuidedFilter(nn.Module):
|
|
def __init__(
|
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
|
):
|
|
super(AdaGuidedFilter, self).__init__()
|
|
|
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
|
self.fc = nn.Conv2d(
|
|
in_channels=n_feats,
|
|
out_channels=1,
|
|
kernel_size=1,
|
|
padding=0,
|
|
stride=1,
|
|
groups=1,
|
|
bias=True,
|
|
)
|
|
|
|
self.r = 5
|
|
|
|
def box_filter(self, x, r):
|
|
channel = x.shape[1]
|
|
kernel_size = 2 * r + 1
|
|
weight = 1.0 / (kernel_size**2)
|
|
box_kernel = weight * torch.ones(
|
|
(channel, 1, kernel_size, kernel_size), dtype=torch.float32, device=x.device
|
|
)
|
|
output = F.conv2d(x, weight=box_kernel, stride=1, padding=r, groups=channel)
|
|
return output
|
|
|
|
def forward(self, x):
|
|
_, _, H, W = x.shape
|
|
N = self.box_filter(
|
|
torch.ones((1, 1, H, W), dtype=x.dtype, device=x.device), self.r
|
|
)
|
|
|
|
# epsilon = self.fc(self.gap(x))
|
|
# epsilon = torch.pow(epsilon, 2)
|
|
epsilon = 1e-2
|
|
|
|
mean_x = self.box_filter(x, self.r) / N
|
|
var_x = self.box_filter(x * x, self.r) / N - mean_x * mean_x
|
|
|
|
A = var_x / (var_x + epsilon)
|
|
b = (1 - A) * mean_x
|
|
m = A * x + b
|
|
|
|
# mean_A = self.box_filter(A, self.r) / N
|
|
# mean_b = self.box_filter(b, self.r) / N
|
|
# m = mean_A * x + mean_b
|
|
return x * m
|
|
|
|
|
|
class AdaConvGuidedFilter(nn.Module):
|
|
def __init__(
|
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
|
):
|
|
super(AdaConvGuidedFilter, self).__init__()
|
|
f = esa_channels
|
|
|
|
self.conv_f = conv(f, f, kernel_size=1)
|
|
|
|
kernel_size = 17
|
|
kernel_expand = kernel_expand
|
|
padding = kernel_size // 2
|
|
|
|
self.vec_conv = nn.Conv2d(
|
|
in_channels=f,
|
|
out_channels=f,
|
|
kernel_size=(1, kernel_size),
|
|
padding=(0, padding),
|
|
groups=f,
|
|
bias=bias,
|
|
)
|
|
|
|
self.hor_conv = nn.Conv2d(
|
|
in_channels=f,
|
|
out_channels=f,
|
|
kernel_size=(kernel_size, 1),
|
|
padding=(padding, 0),
|
|
groups=f,
|
|
bias=bias,
|
|
)
|
|
|
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
|
self.fc = nn.Conv2d(
|
|
in_channels=f,
|
|
out_channels=f,
|
|
kernel_size=1,
|
|
padding=0,
|
|
stride=1,
|
|
groups=1,
|
|
bias=True,
|
|
)
|
|
|
|
def forward(self, x):
|
|
y = self.vec_conv(x)
|
|
y = self.hor_conv(y)
|
|
|
|
sigma = torch.pow(y, 2)
|
|
epsilon = self.fc(self.gap(y))
|
|
|
|
weight = sigma / (sigma + epsilon)
|
|
|
|
m = weight * x + (1 - weight)
|
|
|
|
return x * m
|