Andyx commited on
Commit
b1c6042
·
0 Parent(s):
Files changed (12) hide show
  1. .gitattributes +3 -0
  2. 001.jpg +3 -0
  3. 002.jpg +3 -0
  4. 003.jpg +3 -0
  5. 004.jpg +3 -0
  6. 005.jpg +3 -0
  7. README.md +12 -0
  8. app.py +115 -0
  9. mix.pth +3 -0
  10. model/nets.py +259 -0
  11. requirements.txt +6 -0
  12. 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

  • SHA256: a57fa193ce0c4e532f9ce55169eee2f83475b0422485744fa82a3ae3abfcc3b6
  • Pointer size: 132 Bytes
  • Size of remote file: 5.07 MB
002.jpg ADDED

Git LFS Details

  • SHA256: b9181c80dd860b54fbc8da4dee1565aeffb5c55b24fd2d5a819edf13ba4a2993
  • Pointer size: 132 Bytes
  • Size of remote file: 4.8 MB
003.jpg ADDED

Git LFS Details

  • SHA256: 03d8605cf8d695d775a1c4fe172e6313447cf74a34fa9272f168df13ffb04d09
  • Pointer size: 132 Bytes
  • Size of remote file: 4.84 MB
004.jpg ADDED

Git LFS Details

  • SHA256: 9ac79dfac7ab723dd3f07589aa7ae7e680ee22c6a37800bd79cce6cd82940c35
  • Pointer size: 129 Bytes
  • Size of remote file: 8.12 kB
005.jpg ADDED

Git LFS Details

  • SHA256: 3213c47088e5d0e78a23a1e51abc9d3ce202b16524705d2a7a54cf91a6979819
  • Pointer size: 129 Bytes
  • Size of remote file: 7.75 kB
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