Commit
·
b1c6042
0
Parent(s):
- .gitattributes +3 -0
- 001.jpg +3 -0
- 002.jpg +3 -0
- 003.jpg +3 -0
- 004.jpg +3 -0
- 005.jpg +3 -0
- README.md +12 -0
- app.py +115 -0
- mix.pth +3 -0
- model/nets.py +259 -0
- requirements.txt +6 -0
- uhdm_checkpoint.pth +3 -0
.gitattributes
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
2 |
+
*.jpg filter=lfs diff=lfs merge=lfs -text
|
3 |
+
.jpg filter=lfs diff=lfs merge=lfs -text
|
001.jpg
ADDED
![]() |
Git LFS Details
|
002.jpg
ADDED
![]() |
Git LFS Details
|
003.jpg
ADDED
![]() |
Git LFS Details
|
004.jpg
ADDED
![]() |
Git LFS Details
|
005.jpg
ADDED
![]() |
Git LFS Details
|
README.md
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
title: Screen Image Demoireing
|
3 |
+
emoji: ⚡
|
4 |
+
colorFrom: purple
|
5 |
+
colorTo: purple
|
6 |
+
sdk: gradio
|
7 |
+
sdk_version: 3.1.1
|
8 |
+
app_file: app.py
|
9 |
+
pinned: false
|
10 |
+
---
|
11 |
+
|
12 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
ADDED
@@ -0,0 +1,115 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
from model.nets import my_model
|
3 |
+
import torch
|
4 |
+
import cv2
|
5 |
+
import torch.utils.data as data
|
6 |
+
import torchvision.transforms as transforms
|
7 |
+
import PIL
|
8 |
+
from PIL import Image
|
9 |
+
from PIL import ImageFile
|
10 |
+
import math
|
11 |
+
import os
|
12 |
+
import torch.nn.functional as F
|
13 |
+
|
14 |
+
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
15 |
+
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
16 |
+
model1 = my_model(en_feature_num=48,
|
17 |
+
en_inter_num=32,
|
18 |
+
de_feature_num=64,
|
19 |
+
de_inter_num=32,
|
20 |
+
sam_number=1,
|
21 |
+
).to(device)
|
22 |
+
|
23 |
+
load_path1 = "./mix.pth"
|
24 |
+
model_state_dict1 = torch.load(load_path1, map_location=device)
|
25 |
+
model1.load_state_dict(model_state_dict1)
|
26 |
+
|
27 |
+
model2 = my_model(en_feature_num=48,
|
28 |
+
en_inter_num=32,
|
29 |
+
de_feature_num=64,
|
30 |
+
de_inter_num=32,
|
31 |
+
sam_number=1,
|
32 |
+
).to(device)
|
33 |
+
|
34 |
+
load_path2 = "./uhdm_checkpoint.pth"
|
35 |
+
model_state_dict2 = torch.load(load_path2, map_location=device)
|
36 |
+
model2.load_state_dict(model_state_dict2)
|
37 |
+
|
38 |
+
def default_toTensor(img):
|
39 |
+
t_list = [transforms.ToTensor()]
|
40 |
+
composed_transform = transforms.Compose(t_list)
|
41 |
+
return composed_transform(img)
|
42 |
+
|
43 |
+
def predict1(img):
|
44 |
+
in_img = transforms.ToTensor()(img).to(device).unsqueeze(0)
|
45 |
+
b, c, h, w = in_img.size()
|
46 |
+
# pad image such that the resolution is a multiple of 32
|
47 |
+
w_pad = (math.ceil(w / 32) * 32 - w) // 2
|
48 |
+
h_pad = (math.ceil(h / 32) * 32 - h) // 2
|
49 |
+
in_img = img_pad(in_img, w_r=w_pad, h_r=h_pad)
|
50 |
+
with torch.no_grad():
|
51 |
+
out_1, out_2, out_3 = model1(in_img)
|
52 |
+
if h_pad != 0:
|
53 |
+
out_1 = out_1[:, :, h_pad:-h_pad, :]
|
54 |
+
if w_pad != 0:
|
55 |
+
out_1 = out_1[:, :, :, w_pad:-w_pad]
|
56 |
+
out_1 = out_1.squeeze(0)
|
57 |
+
out_1 = PIL.Image.fromarray(torch.clamp(out_1 * 255, min=0, max=255
|
58 |
+
).byte().permute(1, 2, 0).cpu().numpy())
|
59 |
+
|
60 |
+
return out_1
|
61 |
+
|
62 |
+
def predict2(img):
|
63 |
+
in_img = transforms.ToTensor()(img).to(device).unsqueeze(0)
|
64 |
+
b, c, h, w = in_img.size()
|
65 |
+
# pad image such that the resolution is a multiple of 32
|
66 |
+
w_pad = (math.ceil(w / 32) * 32 - w) // 2
|
67 |
+
h_pad = (math.ceil(h / 32) * 32 - h) // 2
|
68 |
+
in_img = img_pad(in_img, w_r=w_pad, h_r=h_pad)
|
69 |
+
with torch.no_grad():
|
70 |
+
out_1, out_2, out_3 = model2(in_img)
|
71 |
+
if h_pad != 0:
|
72 |
+
out_1 = out_1[:, :, h_pad:-h_pad, :]
|
73 |
+
if w_pad != 0:
|
74 |
+
out_1 = out_1[:, :, :, w_pad:-w_pad]
|
75 |
+
out_1 = out_1.squeeze(0)
|
76 |
+
out_1 = PIL.Image.fromarray(torch.clamp(out_1 * 255, min=0, max=255
|
77 |
+
).byte().permute(1, 2, 0).cpu().numpy())
|
78 |
+
|
79 |
+
return out_1
|
80 |
+
|
81 |
+
def img_pad(x, h_r=0, w_r=0):
|
82 |
+
'''
|
83 |
+
Here the padding values are determined by the average r,g,b values across the training set
|
84 |
+
in FHDMi dataset. For the evaluation on the UHDM, you can also try the commented lines where
|
85 |
+
the mean values are calculated from UHDM training set, yielding similar performance.
|
86 |
+
'''
|
87 |
+
x1 = F.pad(x[:, 0:1, ...], (w_r, w_r, h_r, h_r), value=0.3827)
|
88 |
+
x2 = F.pad(x[:, 1:2, ...], (w_r, w_r, h_r, h_r), value=0.4141)
|
89 |
+
x3 = F.pad(x[:, 2:3, ...], (w_r, w_r, h_r, h_r), value=0.3912)
|
90 |
+
|
91 |
+
y = torch.cat([x1, x2, x3], dim=1)
|
92 |
+
|
93 |
+
return y
|
94 |
+
|
95 |
+
img1 = Image.open('./imgs/001.jpg').convert('RGB')
|
96 |
+
img2 = Image.open('./imgs/002.jpg').convert('RGB')
|
97 |
+
img3 = Image.open('./imgs/003.jpg').convert('RGB')
|
98 |
+
img4 = Image.open('./imgs/004.jpg').convert('RGB')
|
99 |
+
img5 = Image.open('./imgs/005.jpg').convert('RGB')
|
100 |
+
|
101 |
+
|
102 |
+
iface1 = gr.Interface(fn=predict1,
|
103 |
+
inputs=gr.inputs.Image(type="pil"),
|
104 |
+
outputs=gr.inputs.Image(type="pil"))
|
105 |
+
|
106 |
+
iface2 = gr.Interface(fn=predict2,
|
107 |
+
inputs=gr.inputs.Image(type="pil"),
|
108 |
+
outputs=gr.inputs.Image(type="pil"))
|
109 |
+
|
110 |
+
iface_all = gr.mix.Parallel(
|
111 |
+
iface1, iface2,
|
112 |
+
examples=[img1, img2, img3, img4, img5]
|
113 |
+
)
|
114 |
+
|
115 |
+
iface_all.launch()
|
mix.pth
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:bdcdd33f11e1d5eb836671f15991ecb42134bd5ba98c1e4de3b8e2f4138fdb2b
|
3 |
+
size 23895301
|
model/nets.py
ADDED
@@ -0,0 +1,259 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""
|
2 |
+
Implementation of ESDNet for image demoireing
|
3 |
+
"""
|
4 |
+
|
5 |
+
|
6 |
+
import torch
|
7 |
+
import torch.nn as nn
|
8 |
+
import torch.nn.functional as F
|
9 |
+
import torchvision
|
10 |
+
from torch.nn.parameter import Parameter
|
11 |
+
|
12 |
+
class my_model(nn.Module):
|
13 |
+
def __init__(self,
|
14 |
+
en_feature_num,
|
15 |
+
en_inter_num,
|
16 |
+
de_feature_num,
|
17 |
+
de_inter_num,
|
18 |
+
sam_number=1,
|
19 |
+
):
|
20 |
+
super(my_model, self).__init__()
|
21 |
+
self.encoder = Encoder(feature_num=en_feature_num, inter_num=en_inter_num, sam_number=sam_number)
|
22 |
+
self.decoder = Decoder(en_num=en_feature_num, feature_num=de_feature_num, inter_num=de_inter_num,
|
23 |
+
sam_number=sam_number)
|
24 |
+
|
25 |
+
def forward(self, x):
|
26 |
+
y_1, y_2, y_3 = self.encoder(x)
|
27 |
+
out_1, out_2, out_3 = self.decoder(y_1, y_2, y_3)
|
28 |
+
|
29 |
+
return out_1, out_2, out_3
|
30 |
+
|
31 |
+
def _initialize_weights(self):
|
32 |
+
for m in self.modules():
|
33 |
+
if isinstance(m, nn.Conv2d):
|
34 |
+
m.weight.data.normal_(0.0, 0.02)
|
35 |
+
if m.bias is not None:
|
36 |
+
m.bias.data.normal_(0.0, 0.02)
|
37 |
+
if isinstance(m, nn.ConvTranspose2d):
|
38 |
+
m.weight.data.normal_(0.0, 0.02)
|
39 |
+
|
40 |
+
|
41 |
+
class Decoder(nn.Module):
|
42 |
+
def __init__(self, en_num, feature_num, inter_num, sam_number):
|
43 |
+
super(Decoder, self).__init__()
|
44 |
+
self.preconv_3 = conv_relu(4 * en_num, feature_num, 3, padding=1)
|
45 |
+
self.decoder_3 = Decoder_Level(feature_num, inter_num, sam_number)
|
46 |
+
|
47 |
+
self.preconv_2 = conv_relu(2 * en_num + feature_num, feature_num, 3, padding=1)
|
48 |
+
self.decoder_2 = Decoder_Level(feature_num, inter_num, sam_number)
|
49 |
+
|
50 |
+
self.preconv_1 = conv_relu(en_num + feature_num, feature_num, 3, padding=1)
|
51 |
+
self.decoder_1 = Decoder_Level(feature_num, inter_num, sam_number)
|
52 |
+
|
53 |
+
def forward(self, y_1, y_2, y_3):
|
54 |
+
x_3 = y_3
|
55 |
+
x_3 = self.preconv_3(x_3)
|
56 |
+
out_3, feat_3 = self.decoder_3(x_3)
|
57 |
+
|
58 |
+
x_2 = torch.cat([y_2, feat_3], dim=1)
|
59 |
+
x_2 = self.preconv_2(x_2)
|
60 |
+
out_2, feat_2 = self.decoder_2(x_2)
|
61 |
+
|
62 |
+
x_1 = torch.cat([y_1, feat_2], dim=1)
|
63 |
+
x_1 = self.preconv_1(x_1)
|
64 |
+
out_1 = self.decoder_1(x_1, feat=False)
|
65 |
+
|
66 |
+
return out_1, out_2, out_3
|
67 |
+
|
68 |
+
|
69 |
+
class Encoder(nn.Module):
|
70 |
+
def __init__(self, feature_num, inter_num, sam_number):
|
71 |
+
super(Encoder, self).__init__()
|
72 |
+
self.conv_first = nn.Sequential(
|
73 |
+
nn.Conv2d(12, feature_num, kernel_size=5, stride=1, padding=2, bias=True),
|
74 |
+
nn.ReLU(inplace=True)
|
75 |
+
)
|
76 |
+
self.encoder_1 = Encoder_Level(feature_num, inter_num, level=1, sam_number=sam_number)
|
77 |
+
self.encoder_2 = Encoder_Level(2 * feature_num, inter_num, level=2, sam_number=sam_number)
|
78 |
+
self.encoder_3 = Encoder_Level(4 * feature_num, inter_num, level=3, sam_number=sam_number)
|
79 |
+
|
80 |
+
def forward(self, x):
|
81 |
+
x = F.pixel_unshuffle(x, 2)
|
82 |
+
x = self.conv_first(x)
|
83 |
+
|
84 |
+
out_feature_1, down_feature_1 = self.encoder_1(x)
|
85 |
+
out_feature_2, down_feature_2 = self.encoder_2(down_feature_1)
|
86 |
+
out_feature_3 = self.encoder_3(down_feature_2)
|
87 |
+
|
88 |
+
return out_feature_1, out_feature_2, out_feature_3
|
89 |
+
|
90 |
+
|
91 |
+
class Encoder_Level(nn.Module):
|
92 |
+
def __init__(self, feature_num, inter_num, level, sam_number):
|
93 |
+
super(Encoder_Level, self).__init__()
|
94 |
+
self.rdb = RDB(in_channel=feature_num, d_list=(1, 2, 1), inter_num=inter_num)
|
95 |
+
self.sam_blocks = nn.ModuleList()
|
96 |
+
for _ in range(sam_number):
|
97 |
+
sam_block = SAM(in_channel=feature_num, d_list=(1, 2, 3, 2, 1), inter_num=inter_num)
|
98 |
+
self.sam_blocks.append(sam_block)
|
99 |
+
|
100 |
+
if level < 3:
|
101 |
+
self.down = nn.Sequential(
|
102 |
+
nn.Conv2d(feature_num, 2 * feature_num, kernel_size=3, stride=2, padding=1, bias=True),
|
103 |
+
nn.ReLU(inplace=True)
|
104 |
+
)
|
105 |
+
self.level = level
|
106 |
+
|
107 |
+
def forward(self, x):
|
108 |
+
out_feature = self.rdb(x)
|
109 |
+
for sam_block in self.sam_blocks:
|
110 |
+
out_feature = sam_block(out_feature)
|
111 |
+
if self.level < 3:
|
112 |
+
down_feature = self.down(out_feature)
|
113 |
+
return out_feature, down_feature
|
114 |
+
return out_feature
|
115 |
+
|
116 |
+
|
117 |
+
class Decoder_Level(nn.Module):
|
118 |
+
def __init__(self, feature_num, inter_num, sam_number):
|
119 |
+
super(Decoder_Level, self).__init__()
|
120 |
+
self.rdb = RDB(feature_num, (1, 2, 1), inter_num)
|
121 |
+
self.sam_blocks = nn.ModuleList()
|
122 |
+
for _ in range(sam_number):
|
123 |
+
sam_block = SAM(in_channel=feature_num, d_list=(1, 2, 3, 2, 1), inter_num=inter_num)
|
124 |
+
self.sam_blocks.append(sam_block)
|
125 |
+
self.conv = conv(in_channel=feature_num, out_channel=12, kernel_size=3, padding=1)
|
126 |
+
|
127 |
+
def forward(self, x, feat=True):
|
128 |
+
x = self.rdb(x)
|
129 |
+
for sam_block in self.sam_blocks:
|
130 |
+
x = sam_block(x)
|
131 |
+
out = self.conv(x)
|
132 |
+
out = F.pixel_shuffle(out, 2)
|
133 |
+
|
134 |
+
if feat:
|
135 |
+
feature = F.interpolate(x, scale_factor=2, mode='bilinear')
|
136 |
+
return out, feature
|
137 |
+
else:
|
138 |
+
return out
|
139 |
+
|
140 |
+
|
141 |
+
class DB(nn.Module):
|
142 |
+
def __init__(self, in_channel, d_list, inter_num):
|
143 |
+
super(DB, self).__init__()
|
144 |
+
self.d_list = d_list
|
145 |
+
self.conv_layers = nn.ModuleList()
|
146 |
+
c = in_channel
|
147 |
+
for i in range(len(d_list)):
|
148 |
+
dense_conv = conv_relu(in_channel=c, out_channel=inter_num, kernel_size=3, dilation_rate=d_list[i],
|
149 |
+
padding=d_list[i])
|
150 |
+
self.conv_layers.append(dense_conv)
|
151 |
+
c = c + inter_num
|
152 |
+
self.conv_post = conv(in_channel=c, out_channel=in_channel, kernel_size=1)
|
153 |
+
|
154 |
+
def forward(self, x):
|
155 |
+
t = x
|
156 |
+
for conv_layer in self.conv_layers:
|
157 |
+
_t = conv_layer(t)
|
158 |
+
t = torch.cat([_t, t], dim=1)
|
159 |
+
t = self.conv_post(t)
|
160 |
+
return t
|
161 |
+
|
162 |
+
|
163 |
+
class SAM(nn.Module):
|
164 |
+
def __init__(self, in_channel, d_list, inter_num):
|
165 |
+
super(SAM, self).__init__()
|
166 |
+
self.basic_block = DB(in_channel=in_channel, d_list=d_list, inter_num=inter_num)
|
167 |
+
self.basic_block_2 = DB(in_channel=in_channel, d_list=d_list, inter_num=inter_num)
|
168 |
+
self.basic_block_4 = DB(in_channel=in_channel, d_list=d_list, inter_num=inter_num)
|
169 |
+
self.fusion = CSAF(3 * in_channel)
|
170 |
+
|
171 |
+
def forward(self, x):
|
172 |
+
x_0 = x
|
173 |
+
x_2 = F.interpolate(x, scale_factor=0.5, mode='bilinear')
|
174 |
+
x_4 = F.interpolate(x, scale_factor=0.25, mode='bilinear')
|
175 |
+
|
176 |
+
y_0 = self.basic_block(x_0)
|
177 |
+
y_2 = self.basic_block_2(x_2)
|
178 |
+
y_4 = self.basic_block_4(x_4)
|
179 |
+
|
180 |
+
y_2 = F.interpolate(y_2, scale_factor=2, mode='bilinear')
|
181 |
+
y_4 = F.interpolate(y_4, scale_factor=4, mode='bilinear')
|
182 |
+
|
183 |
+
y = self.fusion(y_0, y_2, y_4)
|
184 |
+
y = x + y
|
185 |
+
|
186 |
+
return y
|
187 |
+
|
188 |
+
|
189 |
+
class CSAF(nn.Module):
|
190 |
+
def __init__(self, in_chnls, ratio=4):
|
191 |
+
super(CSAF, self).__init__()
|
192 |
+
self.squeeze = nn.AdaptiveAvgPool2d((1, 1))
|
193 |
+
self.compress1 = nn.Conv2d(in_chnls, in_chnls // ratio, 1, 1, 0)
|
194 |
+
self.compress2 = nn.Conv2d(in_chnls // ratio, in_chnls // ratio, 1, 1, 0)
|
195 |
+
self.excitation = nn.Conv2d(in_chnls // ratio, in_chnls, 1, 1, 0)
|
196 |
+
|
197 |
+
def forward(self, x0, x2, x4):
|
198 |
+
out0 = self.squeeze(x0)
|
199 |
+
out2 = self.squeeze(x2)
|
200 |
+
out4 = self.squeeze(x4)
|
201 |
+
out = torch.cat([out0, out2, out4], dim=1)
|
202 |
+
out = self.compress1(out)
|
203 |
+
out = F.relu(out)
|
204 |
+
out = self.compress2(out)
|
205 |
+
out = F.relu(out)
|
206 |
+
out = self.excitation(out)
|
207 |
+
out = F.sigmoid(out)
|
208 |
+
w0, w2, w4 = torch.chunk(out, 3, dim=1)
|
209 |
+
x = x0 * w0 + x2 * w2 + x4 * w4
|
210 |
+
|
211 |
+
return x
|
212 |
+
|
213 |
+
|
214 |
+
class RDB(nn.Module):
|
215 |
+
def __init__(self, in_channel, d_list, inter_num):
|
216 |
+
super(RDB, self).__init__()
|
217 |
+
self.d_list = d_list
|
218 |
+
self.conv_layers = nn.ModuleList()
|
219 |
+
c = in_channel
|
220 |
+
for i in range(len(d_list)):
|
221 |
+
dense_conv = conv_relu(in_channel=c, out_channel=inter_num, kernel_size=3, dilation_rate=d_list[i],
|
222 |
+
padding=d_list[i])
|
223 |
+
self.conv_layers.append(dense_conv)
|
224 |
+
c = c + inter_num
|
225 |
+
self.conv_post = conv(in_channel=c, out_channel=in_channel, kernel_size=1)
|
226 |
+
|
227 |
+
def forward(self, x):
|
228 |
+
t = x
|
229 |
+
for conv_layer in self.conv_layers:
|
230 |
+
_t = conv_layer(t)
|
231 |
+
t = torch.cat([_t, t], dim=1)
|
232 |
+
|
233 |
+
t = self.conv_post(t)
|
234 |
+
return t + x
|
235 |
+
|
236 |
+
|
237 |
+
class conv(nn.Module):
|
238 |
+
def __init__(self, in_channel, out_channel, kernel_size, dilation_rate=1, padding=0, stride=1):
|
239 |
+
super(conv, self).__init__()
|
240 |
+
self.conv = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=stride,
|
241 |
+
padding=padding, bias=True, dilation=dilation_rate)
|
242 |
+
|
243 |
+
def forward(self, x_input):
|
244 |
+
out = self.conv(x_input)
|
245 |
+
return out
|
246 |
+
|
247 |
+
|
248 |
+
class conv_relu(nn.Module):
|
249 |
+
def __init__(self, in_channel, out_channel, kernel_size, dilation_rate=1, padding=0, stride=1):
|
250 |
+
super(conv_relu, self).__init__()
|
251 |
+
self.conv = nn.Sequential(
|
252 |
+
nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=stride,
|
253 |
+
padding=padding, bias=True, dilation=dilation_rate),
|
254 |
+
nn.ReLU(inplace=True)
|
255 |
+
)
|
256 |
+
|
257 |
+
def forward(self, x_input):
|
258 |
+
out = self.conv(x_input)
|
259 |
+
return out
|
requirements.txt
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
numpy==1.21.5
|
2 |
+
torch>=1.9.0
|
3 |
+
opencv-python==4.5.5.64
|
4 |
+
scikit-image==0.19.2
|
5 |
+
torchvision==0.1.8
|
6 |
+
|
uhdm_checkpoint.pth
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:254235cd25f90a3f1785885385dc6cb3f2178e053291ab53d1943bd7c2f7de65
|
3 |
+
size 23895301
|