Spaces:
Runtime error
Runtime error
Create SRMNet.py
Browse files- model/SRMNet.py +227 -0
model/SRMNet.py
ADDED
@@ -0,0 +1,227 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
|
4 |
+
##---------- Basic Layers ----------
|
5 |
+
def conv3x3(in_chn, out_chn, bias=True):
|
6 |
+
layer = nn.Conv2d(in_chn, out_chn, kernel_size=3, stride=1, padding=1, bias=bias)
|
7 |
+
return layer
|
8 |
+
|
9 |
+
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
|
10 |
+
return nn.Conv2d(
|
11 |
+
in_channels, out_channels, kernel_size,
|
12 |
+
padding=(kernel_size // 2), bias=bias, stride=stride)
|
13 |
+
|
14 |
+
def bili_resize(factor):
|
15 |
+
return nn.Upsample(scale_factor=factor, mode='bilinear', align_corners=False)
|
16 |
+
|
17 |
+
##---------- Basic Blocks ----------
|
18 |
+
class UNetConvBlock(nn.Module):
|
19 |
+
def __init__(self, in_size, out_size, downsample):
|
20 |
+
super(UNetConvBlock, self).__init__()
|
21 |
+
self.downsample = downsample
|
22 |
+
self.block = SK_RDB(in_channels=in_size, growth_rate=out_size, num_layers=3)
|
23 |
+
if downsample:
|
24 |
+
self.downsample = PS_down(out_size, out_size, downscale=2)
|
25 |
+
|
26 |
+
def forward(self, x):
|
27 |
+
out = self.block(x)
|
28 |
+
if self.downsample:
|
29 |
+
out_down = self.downsample(out)
|
30 |
+
return out_down, out
|
31 |
+
else:
|
32 |
+
return out
|
33 |
+
|
34 |
+
class UNetUpBlock(nn.Module):
|
35 |
+
def __init__(self, in_size, out_size):
|
36 |
+
super(UNetUpBlock, self).__init__()
|
37 |
+
# self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
|
38 |
+
self.up = PS_up(in_size, out_size, upscale=2)
|
39 |
+
self.conv_block = UNetConvBlock(in_size, out_size, False)
|
40 |
+
|
41 |
+
def forward(self, x, bridge):
|
42 |
+
up = self.up(x)
|
43 |
+
out = torch.cat([up, bridge], dim=1)
|
44 |
+
out = self.conv_block(out)
|
45 |
+
return out
|
46 |
+
|
47 |
+
##---------- Resizing Modules (Pixel(Un)Shuffle) ----------
|
48 |
+
class PS_down(nn.Module):
|
49 |
+
def __init__(self, in_size, out_size, downscale):
|
50 |
+
super(PS_down, self).__init__()
|
51 |
+
self.UnPS = nn.PixelUnshuffle(downscale)
|
52 |
+
self.conv1 = nn.Conv2d((downscale**2) * in_size, out_size, 1, 1, 0)
|
53 |
+
|
54 |
+
def forward(self, x):
|
55 |
+
x = self.UnPS(x) # h/2, w/2, 4*c
|
56 |
+
x = self.conv1(x)
|
57 |
+
return x
|
58 |
+
|
59 |
+
class PS_up(nn.Module):
|
60 |
+
def __init__(self, in_size, out_size, upscale):
|
61 |
+
super(PS_up, self).__init__()
|
62 |
+
|
63 |
+
self.PS = nn.PixelShuffle(upscale)
|
64 |
+
self.conv1 = nn.Conv2d(in_size//(upscale**2), out_size, 1, 1, 0)
|
65 |
+
|
66 |
+
def forward(self, x):
|
67 |
+
x = self.PS(x) # h/2, w/2, 4*c
|
68 |
+
x = self.conv1(x)
|
69 |
+
return x
|
70 |
+
|
71 |
+
##---------- Selective Kernel Feature Fusion (SKFF) ----------
|
72 |
+
class SKFF(nn.Module):
|
73 |
+
def __init__(self, in_channels, height=3, reduction=8, bias=False):
|
74 |
+
super(SKFF, self).__init__()
|
75 |
+
|
76 |
+
self.height = height
|
77 |
+
d = max(int(in_channels / reduction), 4)
|
78 |
+
|
79 |
+
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
80 |
+
self.conv_du = nn.Sequential(nn.Conv2d(in_channels, d, 1, padding=0, bias=bias), nn.PReLU())
|
81 |
+
|
82 |
+
self.fcs = nn.ModuleList([])
|
83 |
+
for i in range(self.height):
|
84 |
+
self.fcs.append(nn.Conv2d(d, in_channels, kernel_size=1, stride=1, bias=bias))
|
85 |
+
|
86 |
+
self.softmax = nn.Softmax(dim=1)
|
87 |
+
|
88 |
+
def forward(self, inp_feats):
|
89 |
+
batch_size, n_feats, H, W = inp_feats[1].shape
|
90 |
+
|
91 |
+
inp_feats = torch.cat(inp_feats, dim=1)
|
92 |
+
inp_feats = inp_feats.view(batch_size, self.height, n_feats, inp_feats.shape[2], inp_feats.shape[3])
|
93 |
+
|
94 |
+
feats_U = torch.sum(inp_feats, dim=1)
|
95 |
+
feats_S = self.avg_pool(feats_U)
|
96 |
+
feats_Z = self.conv_du(feats_S)
|
97 |
+
|
98 |
+
attention_vectors = [fc(feats_Z) for fc in self.fcs]
|
99 |
+
attention_vectors = torch.cat(attention_vectors, dim=1)
|
100 |
+
attention_vectors = attention_vectors.view(batch_size, self.height, n_feats, 1, 1)
|
101 |
+
|
102 |
+
attention_vectors = self.softmax(attention_vectors)
|
103 |
+
feats_V = torch.sum(inp_feats * attention_vectors, dim=1)
|
104 |
+
|
105 |
+
return feats_V
|
106 |
+
|
107 |
+
##---------- Dense Block ----------
|
108 |
+
class DenseLayer(nn.Module):
|
109 |
+
def __init__(self, in_channels, out_channels, I):
|
110 |
+
super(DenseLayer, self).__init__()
|
111 |
+
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=3 // 2)
|
112 |
+
self.relu = nn.ReLU(inplace=True)
|
113 |
+
self.sk = SKFF(out_channels, height=2, reduction=8, bias=False)
|
114 |
+
|
115 |
+
def forward(self, x):
|
116 |
+
x1 = self.relu(self.conv(x))
|
117 |
+
# output = torch.cat([x, x1], 1) # -> RDB
|
118 |
+
output = self.sk((x, x1))
|
119 |
+
return output
|
120 |
+
|
121 |
+
##---------- Selective Kernel Residual Dense Block (SK-RDB) ----------
|
122 |
+
class SK_RDB(nn.Module):
|
123 |
+
def __init__(self, in_channels, growth_rate, num_layers):
|
124 |
+
super(SK_RDB, self).__init__()
|
125 |
+
self.identity = nn.Conv2d(in_channels, growth_rate, 1, 1, 0)
|
126 |
+
self.layers = nn.Sequential(
|
127 |
+
*[DenseLayer(in_channels, in_channels, I=i) for i in range(num_layers)]
|
128 |
+
)
|
129 |
+
self.lff = nn.Conv2d(in_channels, growth_rate, kernel_size=1)
|
130 |
+
|
131 |
+
def forward(self, x):
|
132 |
+
res = self.identity(x)
|
133 |
+
x = self.layers(x)
|
134 |
+
x = self.lff(x)
|
135 |
+
return res + x
|
136 |
+
|
137 |
+
##---------- testNet ----------
|
138 |
+
class SRMNet(nn.Module):
|
139 |
+
def __init__(self, in_chn=3, wf=96, depth=4):
|
140 |
+
super(SRMNet, self).__init__()
|
141 |
+
self.depth = depth
|
142 |
+
self.down_path = nn.ModuleList()
|
143 |
+
self.bili_down = bili_resize(0.5)
|
144 |
+
self.conv_01 = nn.Conv2d(in_chn, wf, 3, 1, 1)
|
145 |
+
|
146 |
+
# encoder of UNet
|
147 |
+
prev_channels = 0
|
148 |
+
for i in range(depth): # 0,1,2,3
|
149 |
+
downsample = True if (i + 1) < depth else False
|
150 |
+
self.down_path.append(UNetConvBlock(prev_channels + wf, (2 ** i) * wf, downsample))
|
151 |
+
prev_channels = (2 ** i) * wf
|
152 |
+
|
153 |
+
# decoder of UNet
|
154 |
+
self.up_path = nn.ModuleList()
|
155 |
+
self.skip_conv = nn.ModuleList()
|
156 |
+
self.conv_up = nn.ModuleList()
|
157 |
+
self.bottom_conv = nn.Conv2d(prev_channels, wf, 3, 1, 1)
|
158 |
+
self.bottom_up = bili_resize(2 ** (depth-1))
|
159 |
+
|
160 |
+
for i in reversed(range(depth - 1)):
|
161 |
+
self.up_path.append(UNetUpBlock(prev_channels, (2 ** i) * wf))
|
162 |
+
self.skip_conv.append(nn.Conv2d((2 ** i) * wf, (2 ** i) * wf, 3, 1, 1))
|
163 |
+
self.conv_up.append(nn.Sequential(*[nn.Conv2d((2 ** i) * wf, wf, 3, 1, 1), bili_resize(2 ** i)]))
|
164 |
+
prev_channels = (2 ** i) * wf
|
165 |
+
|
166 |
+
self.final_ff = SKFF(in_channels=wf, height=depth)
|
167 |
+
self.last = conv3x3(prev_channels, in_chn, bias=True)
|
168 |
+
|
169 |
+
def forward(self, x):
|
170 |
+
img = x
|
171 |
+
scale_img = img
|
172 |
+
|
173 |
+
##### shallow conv #####
|
174 |
+
x1 = self.conv_01(img)
|
175 |
+
encs = []
|
176 |
+
######## UNet ########
|
177 |
+
# Down-path (Encoder)
|
178 |
+
for i, down in enumerate(self.down_path):
|
179 |
+
if i == 0:
|
180 |
+
x1, x1_up = down(x1)
|
181 |
+
encs.append(x1_up)
|
182 |
+
elif (i + 1) < self.depth:
|
183 |
+
scale_img = self.bili_down(scale_img)
|
184 |
+
left_bar = self.conv_01(scale_img)
|
185 |
+
x1 = torch.cat([x1, left_bar], dim=1)
|
186 |
+
x1, x1_up = down(x1)
|
187 |
+
encs.append(x1_up)
|
188 |
+
else:
|
189 |
+
scale_img = self.bili_down(scale_img)
|
190 |
+
left_bar = self.conv_01(scale_img)
|
191 |
+
x1 = torch.cat([x1, left_bar], dim=1)
|
192 |
+
x1 = down(x1)
|
193 |
+
|
194 |
+
# Up-path (Decoder)
|
195 |
+
ms_result = [self.bottom_up(self.bottom_conv(x1))]
|
196 |
+
for i, up in enumerate(self.up_path):
|
197 |
+
x1 = up(x1, self.skip_conv[i](encs[-i - 1]))
|
198 |
+
ms_result.append(self.conv_up[i](x1))
|
199 |
+
|
200 |
+
# Multi-scale selective feature fusion
|
201 |
+
msff_result = self.final_ff(ms_result)
|
202 |
+
|
203 |
+
##### Reconstruct #####
|
204 |
+
out_1 = self.last(msff_result) + img
|
205 |
+
|
206 |
+
return out_1
|
207 |
+
|
208 |
+
|
209 |
+
if __name__ == "__main__":
|
210 |
+
from thop import profile
|
211 |
+
|
212 |
+
input = torch.ones(1, 3, 256, 256, dtype=torch.float, requires_grad=False)
|
213 |
+
model = SRMNet(in_chn=3, wf=96, depth=4)
|
214 |
+
out = model(input)
|
215 |
+
flops, params = profile(model, inputs=(input,))
|
216 |
+
total = sum(p.numel() for p in model.parameters())
|
217 |
+
|
218 |
+
# RDBlayer = SK_RDB(in_channels=64, growth_rate=64, num_layers=3)
|
219 |
+
# print(RDBlayer)
|
220 |
+
# out = RDBlayer(input)
|
221 |
+
# flops, params = profile(RDBlayer, inputs=(input,))
|
222 |
+
|
223 |
+
print('input shape:', input.shape)
|
224 |
+
print('output shape', out.shape)
|
225 |
+
print("-----------------------------------")
|
226 |
+
print("Total params: %.4f M" % (total / 1e6))
|
227 |
+
print("Total params: %.4f G" % (flops / 1e9))
|