add convir
This commit is contained in:
FuouM
@@ -0,0 +1,196 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .layers import BasicConv, ResBlock
|
||||
|
||||
|
||||
class EBlock(nn.Module):
|
||||
def __init__(self, out_channel, num_res, data):
|
||||
super(EBlock, self).__init__()
|
||||
|
||||
layers = [ResBlock(out_channel, out_channel, data) for _ in range(num_res - 1)]
|
||||
layers.append(ResBlock(out_channel, out_channel, data, filter=True))
|
||||
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class DBlock(nn.Module):
|
||||
def __init__(self, channel, num_res, data):
|
||||
super(DBlock, self).__init__()
|
||||
|
||||
layers = [ResBlock(channel, channel, data) for _ in range(num_res - 1)]
|
||||
layers.append(ResBlock(channel, channel, data, filter=True))
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class SCM(nn.Module):
|
||||
def __init__(self, out_plane):
|
||||
super(SCM, self).__init__()
|
||||
self.main = nn.Sequential(
|
||||
BasicConv(3, out_plane // 4, kernel_size=3, stride=1, relu=True),
|
||||
BasicConv(
|
||||
out_plane // 4, out_plane // 2, kernel_size=1, stride=1, relu=True
|
||||
),
|
||||
BasicConv(
|
||||
out_plane // 2, out_plane // 2, kernel_size=3, stride=1, relu=True
|
||||
),
|
||||
BasicConv(out_plane // 2, out_plane, kernel_size=1, stride=1, relu=False),
|
||||
nn.InstanceNorm2d(out_plane, affine=True),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.main(x)
|
||||
return x
|
||||
|
||||
|
||||
class FAM(nn.Module):
|
||||
def __init__(self, channel):
|
||||
super(FAM, self).__init__()
|
||||
self.merge = BasicConv(
|
||||
channel * 2, channel, kernel_size=3, stride=1, relu=False
|
||||
)
|
||||
|
||||
def forward(self, x1, x2):
|
||||
return self.merge(torch.cat([x1, x2], dim=1))
|
||||
|
||||
|
||||
class ConvIR(nn.Module):
|
||||
def __init__(self, version, data):
|
||||
super(ConvIR, self).__init__()
|
||||
|
||||
if version == "small":
|
||||
num_res = 4
|
||||
elif version == "base":
|
||||
num_res = 8
|
||||
elif version == "large":
|
||||
num_res = 16
|
||||
|
||||
base_channel = 32
|
||||
|
||||
self.Encoder = nn.ModuleList(
|
||||
[
|
||||
EBlock(base_channel, num_res, data),
|
||||
EBlock(base_channel * 2, num_res, data),
|
||||
EBlock(base_channel * 4, num_res, data),
|
||||
]
|
||||
)
|
||||
|
||||
self.feat_extract = nn.ModuleList(
|
||||
[
|
||||
BasicConv(3, base_channel, kernel_size=3, relu=True, stride=1),
|
||||
BasicConv(
|
||||
base_channel, base_channel * 2, kernel_size=3, relu=True, stride=2
|
||||
),
|
||||
BasicConv(
|
||||
base_channel * 2,
|
||||
base_channel * 4,
|
||||
kernel_size=3,
|
||||
relu=True,
|
||||
stride=2,
|
||||
),
|
||||
BasicConv(
|
||||
base_channel * 4,
|
||||
base_channel * 2,
|
||||
kernel_size=4,
|
||||
relu=True,
|
||||
stride=2,
|
||||
transpose=True,
|
||||
),
|
||||
BasicConv(
|
||||
base_channel * 2,
|
||||
base_channel,
|
||||
kernel_size=4,
|
||||
relu=True,
|
||||
stride=2,
|
||||
transpose=True,
|
||||
),
|
||||
BasicConv(base_channel, 3, kernel_size=3, relu=False, stride=1),
|
||||
]
|
||||
)
|
||||
|
||||
self.Decoder = nn.ModuleList(
|
||||
[
|
||||
DBlock(base_channel * 4, num_res, data),
|
||||
DBlock(base_channel * 2, num_res, data),
|
||||
DBlock(base_channel, num_res, data),
|
||||
]
|
||||
)
|
||||
|
||||
self.Convs = nn.ModuleList(
|
||||
[
|
||||
BasicConv(
|
||||
base_channel * 4,
|
||||
base_channel * 2,
|
||||
kernel_size=1,
|
||||
relu=True,
|
||||
stride=1,
|
||||
),
|
||||
BasicConv(
|
||||
base_channel * 2, base_channel, kernel_size=1, relu=True, stride=1
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
self.ConvsOut = nn.ModuleList(
|
||||
[
|
||||
BasicConv(base_channel * 4, 3, kernel_size=3, relu=False, stride=1),
|
||||
BasicConv(base_channel * 2, 3, kernel_size=3, relu=False, stride=1),
|
||||
]
|
||||
)
|
||||
|
||||
self.FAM1 = FAM(base_channel * 4)
|
||||
self.SCM1 = SCM(base_channel * 4)
|
||||
self.FAM2 = FAM(base_channel * 2)
|
||||
self.SCM2 = SCM(base_channel * 2)
|
||||
|
||||
def forward(self, x):
|
||||
x_2 = F.interpolate(x, scale_factor=0.5)
|
||||
x_4 = F.interpolate(x_2, scale_factor=0.5)
|
||||
z2 = self.SCM2(x_2)
|
||||
z4 = self.SCM1(x_4)
|
||||
|
||||
outputs = list()
|
||||
# 256
|
||||
x_ = self.feat_extract[0](x)
|
||||
res1 = self.Encoder[0](x_)
|
||||
# 128
|
||||
z = self.feat_extract[1](res1)
|
||||
z = self.FAM2(z, z2)
|
||||
res2 = self.Encoder[1](z)
|
||||
# 64
|
||||
z = self.feat_extract[2](res2)
|
||||
z = self.FAM1(z, z4)
|
||||
z = self.Encoder[2](z)
|
||||
|
||||
z = self.Decoder[0](z)
|
||||
z_ = self.ConvsOut[0](z)
|
||||
# 128
|
||||
z = self.feat_extract[3](z)
|
||||
outputs.append(z_ + x_4)
|
||||
|
||||
z = torch.cat([z, res2], dim=1)
|
||||
z = self.Convs[0](z)
|
||||
z = self.Decoder[1](z)
|
||||
z_ = self.ConvsOut[1](z)
|
||||
# 256
|
||||
z = self.feat_extract[4](z)
|
||||
outputs.append(z_ + x_2)
|
||||
|
||||
z = torch.cat([z, res1], dim=1)
|
||||
z = self.Convs[1](z)
|
||||
z = self.Decoder[2](z)
|
||||
z = self.feat_extract[5](z)
|
||||
outputs.append(z + x)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def build_net(version, data):
|
||||
return ConvIR(version, data)
|
||||
@@ -0,0 +1,66 @@
|
||||
MODEL_DIR = "models/convir"
|
||||
SIZE_BASE = "base"
|
||||
SIZE_SMALL = "small"
|
||||
SIZE_LARGE = "large"
|
||||
|
||||
########### Image Dehazing ###########
|
||||
DEHAZE_FULL = "haze4k"
|
||||
DEHAZE_MODELS = [
|
||||
"its",
|
||||
"ots",
|
||||
"densehaze",
|
||||
"haze4k",
|
||||
"ihaze",
|
||||
"nhhaze",
|
||||
"ohaze",
|
||||
"gta5",
|
||||
"nhr",
|
||||
]
|
||||
|
||||
DEHAZE_DEFAULT = "haze4k"
|
||||
|
||||
DEHAZE_SIZES = ["small", "base", "large"]
|
||||
|
||||
|
||||
def get_dehaze_model(model_name: str, model_size: str):
|
||||
if model_name != DEHAZE_FULL and model_size == "large":
|
||||
print(f"Invalid size for model {model_name}. Revert to {SIZE_BASE}")
|
||||
model_size = SIZE_BASE
|
||||
return f"dehaze/{model_name}-{model_size}.pkl", model_size
|
||||
|
||||
|
||||
########### Image Deraining ###########
|
||||
DERAIN_MODEL = "deraining"
|
||||
DERAIN_SIZE = "large"
|
||||
|
||||
|
||||
def get_derain_model():
|
||||
return f"{DERAIN_MODEL}.pkl"
|
||||
|
||||
|
||||
########### Image Desnowing ###########
|
||||
DESNOW_MODELS = ["csd", "snow100k", "srrs"]
|
||||
DESNOW_DEFAULT = "snow100k"
|
||||
DESNOW_SIZES = ["small", "base"]
|
||||
|
||||
|
||||
def get_desnow_model(model_name: str, model_size: str):
|
||||
return f"desnow/{model_name}-{model_size}.pkl"
|
||||
|
||||
|
||||
########### Defocus Deblurring ###########
|
||||
DEFOCUS_MODEL = "dpdd"
|
||||
DEFOCUS_SIZES = ["small", "base", "large"]
|
||||
|
||||
|
||||
def get_defocus_model(model_size: str):
|
||||
return f"defocus/{DEFOCUS_MODEL}-{model_size}.pkl"
|
||||
|
||||
|
||||
########### Motion Deblurring ###########
|
||||
DEBLUR_MODELS = ["gopro", "rsblur"]
|
||||
DEBLUR_DEFAULT = "gopro"
|
||||
|
||||
|
||||
def get_deblur_model(model_name: str):
|
||||
return f"modeblur/convir_{model_name}.pkl"
|
||||
@@ -0,0 +1,30 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def preprocess_image(image: torch.Tensor, device, factor=32, min_size=225):
|
||||
# [B, H, W, C]
|
||||
input_img = image.permute(0, 3, 1, 2).to(device) # [B, C, H, W]
|
||||
_, _, h, w = input_img.shape
|
||||
|
||||
# Resize if either dimension is smaller than min_size
|
||||
if h < min_size or w < min_size:
|
||||
scale = min_size / min(h, w)
|
||||
new_h, new_w = int(h * scale), int(w * scale)
|
||||
input_img = F.interpolate(
|
||||
input_img, size=(new_h, new_w), mode="bilinear", align_corners=False
|
||||
)
|
||||
h, w = new_h, new_w
|
||||
|
||||
# Calculate padding
|
||||
H = ((h + factor - 1) // factor) * factor
|
||||
W = ((w + factor - 1) // factor) * factor
|
||||
|
||||
padh = H - h
|
||||
padw = W - w
|
||||
|
||||
# Pad the image
|
||||
if padh > 0 or padw > 0:
|
||||
input_img = F.pad(input_img, (0, padw, 0, padh), mode="reflect")
|
||||
|
||||
return input_img, h, w
|
||||
@@ -0,0 +1,271 @@
|
||||
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
|
||||
Reference in New Issue
Block a user