1

app.py

@@ -24,16 +24,6 @@ load_path1 = "./mix.pth"
 model_state_dict1 = torch.load(load_path1, map_location=device)
 model1.load_state_dict(model_state_dict1)
 
-model2 = my_model(en_feature_num=48,
-                     en_inter_num=32,
-                     de_feature_num=64,
-                     de_inter_num=32,
-                     sam_number=1,
-                     ).to(device)
-
-load_path2 = "./uhdm_checkpoint.pth"
-model_state_dict2 = torch.load(load_path2, map_location=device)
-model2.load_state_dict(model_state_dict2)
 
 def default_toTensor(img):
     t_list = [transforms.ToTensor()]
@@ -59,25 +49,6 @@ def predict1(img):
 
     return out_1
 
-def predict2(img):
-    in_img = transforms.ToTensor()(img).to(device).unsqueeze(0)
-    b, c, h, w = in_img.size()
-    # pad image such that the resolution is a multiple of 32
-    w_pad = (math.ceil(w / 32) * 32 - w) // 2
-    h_pad = (math.ceil(h / 32) * 32 - h) // 2
-    in_img = img_pad(in_img, w_r=w_pad, h_r=h_pad)
-    with torch.no_grad():
-        out_1, out_2, out_3 = model2(in_img)
-        if h_pad != 0:
-            out_1 = out_1[:, :, h_pad:-h_pad, :]
-        if w_pad != 0:
-            out_1 = out_1[:, :, :, w_pad:-w_pad]
-    out_1 = out_1.squeeze(0)
-    out_1 = PIL.Image.fromarray(torch.clamp(out_1 * 255, min=0, max=255
-    ).byte().permute(1, 2, 0).cpu().numpy())
-
-    return out_1
-
 def img_pad(x, h_r=0, w_r=0):
     '''
     Here the padding values are determined by the average r,g,b values across the training set
@@ -93,21 +64,33 @@ def img_pad(x, h_r=0, w_r=0):
     return y
 
 
+title = "Clean Your Moire Images!"
+description = """
+
+The model was trained to remove the moire patterns from your captured screen images! Specially, this model is capable of tackling
+images up to 4K resolution, which adapts to most of the modern mobile phones. 
+(Note: It may cost 80s per 4K image (e.g., iPhone's resolution: 4032x3024) since this demo runs on the CPU. The model can run
+on a NVIDIA 3090 GPU 17ms per standard 4K image)
+
+
+"""
+
+article = "Check out the [ECCV 2022 paper](https://arxiv.org/abs/2207.09935) and the \
+            [official training code](https://github.com/CVMI-Lab/UHDM) which the demo is based on."
+
+
 iface1 = gr.Interface(fn=predict1,
-             inputs=gr.inputs.Image(type="pil"),
-             outputs=gr.inputs.Image(type="pil"))
-
-iface2 = gr.Interface(fn=predict2,
-             inputs=gr.inputs.Image(type="pil"),
-             outputs=gr.inputs.Image(type="pil"))
-
-iface_all = gr.mix.Parallel(
-    iface1, iface2,
-    examples=['001.jpg',
-              '002.jpg',
-              '003.jpg',
-              '004.jpg',
-              '005.jpg']
-)
-
-iface_all.launch()
+                     inputs=gr.inputs.Image(type="pil"),
+                     outputs=gr.inputs.Image(type="pil"),
+                     examples=['001.jpg',
+                            '002.jpg',
+                            '003.jpg',
+                            '004.jpg',
+                            '005.jpg'],
+                    title = title, 
+                    description = description,
+                    article = article
+                    )
+
+
+iface1.launch()

uhdm_checkpoint.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:254235cd25f90a3f1785885385dc6cb3f2178e053291ab53d1943bd7c2f7de65
-size 23895301