-
Notifications
You must be signed in to change notification settings - Fork 0
/
model_temporal.py
137 lines (119 loc) · 3.75 KB
/
model_temporal.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
from package_core.net_basics import conv2d, Cascade_resnet_blocks
from detectron2.layers import ModulatedDeformConv
from network_swinir import SwinIR_NIR_Encoder
def actFunc(act, *args, **kwargs):
act = act.lower()
if act == 'relu':
return nn.ReLU()
elif act == 'relu6':
return nn.ReLU6()
elif act == 'leakyrelu':
return nn.LeakyReLU(0.1)
elif act == 'prelu':
return nn.PReLU()
elif act == 'rrelu':
return nn.RReLU(0.1, 0.3)
elif act == 'selu':
return nn.SELU()
elif act == 'celu':
return nn.CELU()
elif act == 'elu':
return nn.ELU()
elif act == 'gelu':
return nn.GELU()
elif act == 'tanh':
return nn.Tanh()
else:
raise NotImplementedError
def conv3x3(in_channels, out_channels, stride=1):
return nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True)
class ImageEncoder(nn.Module):
def __init__(self, in_chs, init_chs):
super(ImageEncoder, self).__init__()
self.conv0 = conv2d(
in_planes=in_chs,
out_planes=init_chs,
batch_norm=False,
activation=nn.ReLU(),
kernel_size=7,
stride=1
)
self.resblocks0 = Cascade_resnet_blocks(in_planes=init_chs, n_blocks=1) # green block in paper
self.conv1 = conv2d(
in_planes=init_chs,
out_planes=2 * init_chs,
batch_norm=False,
activation=nn.ReLU(),
kernel_size=3,
stride=2
)
self.resblocks1 = Cascade_resnet_blocks(in_planes=2 * init_chs, n_blocks=1)
self.conv2 = conv2d(
in_planes=2 * init_chs,
out_planes=2 * init_chs,
batch_norm=False,
activation=nn.ReLU(),
kernel_size=3,
stride=2
)
self.resblocks2 = Cascade_resnet_blocks(in_planes=2 * init_chs, n_blocks=1)
def forward(self, x):
x0 = self.resblocks0(self.conv0(x))
x1 = self.resblocks1(self.conv1(x0))
x2 = self.resblocks2(self.conv2(x1))
return x2, x1, x0
class ModulatedDeformLayer(nn.Module):
"""
Modulated Deformable Convolution (v2)
"""
def __init__(self, in_chs, out_chs, kernel_size=3, deformable_groups=1, activation='relu'):
super(ModulatedDeformLayer, self).__init__()
assert isinstance(kernel_size, (int, list, tuple))
self.deform_offset = conv3x3(in_chs, (3 * kernel_size ** 2) * deformable_groups)
self.act = actFunc(activation)
self.deform = ModulatedDeformConv(
in_chs,
out_chs,
kernel_size,
stride=1,
padding=1,
deformable_groups=deformable_groups
)
def forward(self, x, feat):
offset_mask = self.deform_offset(feat)
offset_x, offset_y, mask = torch.chunk(offset_mask, 3, dim=1)
offset = torch.cat((offset_x, offset_y), dim=1)
mask = mask.sigmoid()
out = self.deform(x, offset, mask)
out = self.act(out)
return out
class Generator_noflow(nn.Module):
def __init__(self, opts, in_chs=3, out_chs=3, n_feats=64, embed_dim=48):
super(Generator_noflow, self).__init__()
self.swin = SwinIR_NIR_Encoder(
img_size=(240, 320), window_size=8, img_range=1.,
depths=[6, 6], embed_dim=embed_dim,
num_heads=[6, 6], mlp_ratio=2
)
self.deform_0 = ModulatedDeformLayer(3 * embed_dim, n_feats, deformable_groups=8)
self.resblocks_0 = Cascade_resnet_blocks(n_feats, n_blocks=3)
self.deform_1 = ModulatedDeformLayer(n_feats, n_feats, deformable_groups=8)
self.resblocks_1 = Cascade_resnet_blocks(n_feats, n_blocks=3)
self.to_RGB = nn.Conv2d(n_feats, out_chs, kernel_size=3, stride=1, padding=1)
def forward(self, nir_0, nir_1, nir_2, isp_0, isp_1, isp_2):
fea0 = self.swin(nir_0, isp_0)
fea1 = self.swin(nir_1, isp_1)
fea2 = self.swin(nir_2, isp_2)
x = torch.cat([fea0, fea1, fea2], dim=1)
x = self.deform_0(x, x)
x = self.resblocks_0(x)
x = self.deform_1(x, x)
x = self.resblocks_1(x)
x = self.to_RGB(x)
return x, [x, x, x, x], [x, x, x, x]