272 lines
8.6 KiB
Python
272 lines
8.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class BasicConv(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channel,
|
|
out_channel,
|
|
kernel_size,
|
|
stride,
|
|
bias=True,
|
|
norm=False,
|
|
relu=True,
|
|
transpose=False,
|
|
):
|
|
super(BasicConv, self).__init__()
|
|
if bias and norm:
|
|
bias = False
|
|
|
|
padding = kernel_size // 2
|
|
layers = list()
|
|
if transpose:
|
|
padding = kernel_size // 2 - 1
|
|
layers.append(
|
|
nn.ConvTranspose2d(
|
|
in_channel,
|
|
out_channel,
|
|
kernel_size,
|
|
padding=padding,
|
|
stride=stride,
|
|
bias=bias,
|
|
)
|
|
)
|
|
else:
|
|
layers.append(
|
|
nn.Conv2d(
|
|
in_channel,
|
|
out_channel,
|
|
kernel_size,
|
|
padding=padding,
|
|
stride=stride,
|
|
bias=bias,
|
|
)
|
|
)
|
|
if norm:
|
|
layers.append(nn.BatchNorm2d(out_channel))
|
|
if relu:
|
|
layers.append(nn.GELU())
|
|
self.main = nn.Sequential(*layers)
|
|
|
|
def forward(self, x):
|
|
return self.main(x)
|
|
|
|
|
|
class ResBlock(nn.Module):
|
|
def __init__(self, in_channel, out_channel, data, filter=False):
|
|
super(ResBlock, self).__init__()
|
|
self.main = nn.Sequential(
|
|
BasicConv(in_channel, out_channel, kernel_size=3, stride=1, relu=True),
|
|
DeepPoolLayer(in_channel, out_channel, data) if filter else nn.Identity(),
|
|
BasicConv(out_channel, out_channel, kernel_size=3, stride=1, relu=False),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.main(x) + x
|
|
|
|
|
|
class DeepPoolLayer(nn.Module):
|
|
def __init__(self, k, k_out, data):
|
|
super(DeepPoolLayer, self).__init__()
|
|
self.pools_sizes = [8, 4, 2]
|
|
|
|
if data in [
|
|
# DeHaze
|
|
"its",
|
|
"ots",
|
|
"densehaze",
|
|
"haze4k",
|
|
"ihaze",
|
|
"nhhaze",
|
|
"nhr",
|
|
"ohaze",
|
|
# DeSnow
|
|
"csd",
|
|
"snow100k",
|
|
"srrs",
|
|
# DefocusDeBlur
|
|
"dpdd",
|
|
# DeBlur
|
|
"gopro",
|
|
"rsblur",
|
|
]:
|
|
dilation = [7, 9, 11]
|
|
elif data in ["gta5"]:
|
|
dilation = [5, 9, 11]
|
|
elif data in ["deraining"]:
|
|
dilation = [3, 7, 9]
|
|
|
|
pools, convs, dynas = [], [], []
|
|
for j, i in enumerate(self.pools_sizes):
|
|
pools.append(nn.AvgPool2d(kernel_size=i, stride=i))
|
|
convs.append(nn.Conv2d(k, k, 3, 1, 1, bias=False))
|
|
dynas.append(MultiShapeKernel(dim=k, kernel_size=3, dilation=dilation[j]))
|
|
self.pools = nn.ModuleList(pools)
|
|
self.convs = nn.ModuleList(convs)
|
|
self.dynas = nn.ModuleList(dynas)
|
|
self.relu = nn.GELU()
|
|
self.conv_sum = nn.Conv2d(k, k_out, 3, 1, 1, bias=False)
|
|
|
|
def forward(self, x):
|
|
x_size = x.size()
|
|
resl = x
|
|
y_up = None
|
|
for i in range(len(self.pools_sizes)):
|
|
if i == 0:
|
|
y = self.dynas[i](self.convs[i](self.pools[i](x)))
|
|
else:
|
|
y = self.dynas[i](self.convs[i](self.pools[i](x) + y_up))
|
|
resl = torch.add(
|
|
resl, F.interpolate(y, x_size[2:], mode="bilinear", align_corners=True)
|
|
)
|
|
if i != len(self.pools_sizes) - 1:
|
|
y_up = F.interpolate(
|
|
y, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
resl = self.relu(resl)
|
|
resl = self.conv_sum(resl)
|
|
|
|
return resl
|
|
|
|
|
|
class dynamic_filter(nn.Module):
|
|
def __init__(self, inchannels, kernel_size=3, dilation=1, stride=1, group=8):
|
|
super(dynamic_filter, self).__init__()
|
|
self.stride = stride
|
|
self.kernel_size = kernel_size
|
|
self.group = group
|
|
self.dilation = dilation
|
|
|
|
self.conv = nn.Conv2d(
|
|
inchannels, group * kernel_size**2, kernel_size=1, stride=1, bias=False
|
|
)
|
|
self.bn = nn.BatchNorm2d(group * kernel_size**2)
|
|
self.act = nn.Tanh()
|
|
|
|
nn.init.kaiming_normal_(self.conv.weight, mode="fan_out", nonlinearity="relu")
|
|
self.lamb_l = nn.Parameter(torch.zeros(inchannels), requires_grad=True)
|
|
self.lamb_h = nn.Parameter(torch.zeros(inchannels), requires_grad=True)
|
|
self.pad = nn.ReflectionPad2d(self.dilation * (kernel_size - 1) // 2)
|
|
|
|
self.ap = nn.AdaptiveAvgPool2d((1, 1))
|
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
|
|
|
self.inside_all = nn.Parameter(
|
|
torch.zeros(inchannels, 1, 1), requires_grad=True
|
|
)
|
|
|
|
def forward(self, x):
|
|
identity_input = x
|
|
low_filter = self.ap(x)
|
|
low_filter = self.conv(low_filter)
|
|
low_filter = self.bn(low_filter)
|
|
|
|
n, c, h, w = x.shape
|
|
x = F.unfold(
|
|
self.pad(x), kernel_size=self.kernel_size, dilation=self.dilation
|
|
).reshape(n, self.group, c // self.group, self.kernel_size**2, h * w)
|
|
|
|
n, c1, p, q = low_filter.shape
|
|
low_filter = low_filter.reshape(
|
|
n, c1 // self.kernel_size**2, self.kernel_size**2, p * q
|
|
).unsqueeze(2)
|
|
|
|
low_filter = self.act(low_filter)
|
|
|
|
low_part = torch.sum(x * low_filter, dim=3).reshape(n, c, h, w)
|
|
|
|
out_low = low_part * (self.inside_all + 1.0) - self.inside_all * self.gap(
|
|
identity_input
|
|
)
|
|
|
|
out_low = out_low * self.lamb_l[None, :, None, None]
|
|
|
|
out_high = (identity_input) * (self.lamb_h[None, :, None, None] + 1.0)
|
|
|
|
return out_low + out_high
|
|
|
|
|
|
class cubic_attention(nn.Module):
|
|
def __init__(self, dim, group, dilation, kernel) -> None:
|
|
super().__init__()
|
|
|
|
self.H_spatial_att = spatial_strip_att(
|
|
dim, dilation=dilation, group=group, kernel=kernel
|
|
)
|
|
self.W_spatial_att = spatial_strip_att(
|
|
dim, dilation=dilation, group=group, kernel=kernel, H=False
|
|
)
|
|
self.gamma = nn.Parameter(torch.zeros(dim, 1, 1))
|
|
self.beta = nn.Parameter(torch.ones(dim, 1, 1))
|
|
|
|
def forward(self, x):
|
|
out = self.H_spatial_att(x)
|
|
out = self.W_spatial_att(out)
|
|
return self.gamma * out + x * self.beta
|
|
|
|
|
|
class spatial_strip_att(nn.Module):
|
|
def __init__(self, dim, kernel=3, dilation=1, group=2, H=True) -> None:
|
|
super().__init__()
|
|
|
|
self.k = kernel
|
|
pad = dilation * (kernel - 1) // 2
|
|
self.kernel = (1, kernel) if H else (kernel, 1)
|
|
self.padding = (kernel // 2, 1) if H else (1, kernel // 2)
|
|
self.dilation = dilation
|
|
self.group = group
|
|
self.pad = (
|
|
nn.ReflectionPad2d((pad, pad, 0, 0))
|
|
if H
|
|
else nn.ReflectionPad2d((0, 0, pad, pad))
|
|
)
|
|
self.conv = nn.Conv2d(dim, group * kernel, kernel_size=1, stride=1, bias=False)
|
|
self.ap = nn.AdaptiveAvgPool2d((1, 1))
|
|
self.filter_act = nn.Tanh()
|
|
self.inside_all = nn.Parameter(torch.zeros(dim, 1, 1), requires_grad=True)
|
|
self.lamb_l = nn.Parameter(torch.zeros(dim), requires_grad=True)
|
|
self.lamb_h = nn.Parameter(torch.zeros(dim), requires_grad=True)
|
|
gap_kernel = (None, 1) if H else (1, None)
|
|
self.gap = nn.AdaptiveAvgPool2d(gap_kernel)
|
|
|
|
def forward(self, x):
|
|
identity_input = x.clone()
|
|
filter = self.ap(x)
|
|
filter = self.conv(filter)
|
|
n, c, h, w = x.shape
|
|
x = F.unfold(
|
|
self.pad(x), kernel_size=self.kernel, dilation=self.dilation
|
|
).reshape(n, self.group, c // self.group, self.k, h * w)
|
|
n, c1, p, q = filter.shape
|
|
filter = filter.reshape(n, c1 // self.k, self.k, p * q).unsqueeze(2)
|
|
filter = self.filter_act(filter)
|
|
out = torch.sum(x * filter, dim=3).reshape(n, c, h, w)
|
|
|
|
out_low = out * (self.inside_all + 1.0) - self.inside_all * self.gap(
|
|
identity_input
|
|
)
|
|
out_low = out_low * self.lamb_l[None, :, None, None]
|
|
out_high = identity_input * (self.lamb_h[None, :, None, None] + 1.0)
|
|
|
|
return out_low + out_high
|
|
|
|
|
|
class MultiShapeKernel(nn.Module):
|
|
def __init__(self, dim, kernel_size=3, dilation=1, group=8):
|
|
super().__init__()
|
|
|
|
self.square_att = dynamic_filter(
|
|
inchannels=dim, dilation=dilation, group=group, kernel_size=kernel_size
|
|
)
|
|
self.strip_att = cubic_attention(
|
|
dim, group=group, dilation=dilation, kernel=kernel_size
|
|
)
|
|
|
|
def forward(self, x):
|
|
x1 = self.strip_att(x)
|
|
x2 = self.square_att(x)
|
|
|
|
return x1 + x2
|