Upload app.py

app.py

@@ -0,0 +1,116 @@
+import gradio as gr
+import os
+import numpy as np
+import torch
+from models.network_swinir import SwinIR
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("device: %s" % device)
+default_models = {
+    "sr": "weights/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth",
+    "denoise": "weights/005_colorDN_DFWB_s128w8_SwinIR-M_noise25.pth"
+    }
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.benchmark = True
+
+
+denoise_model = SwinIR(upscale=1, in_chans=3, img_size=128, window_size=8,
+        img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
+        mlp_ratio=2, upsampler='', resi_connection='1conv').to(device)
+param_key_g = 'params'
+try:
+    pretrained_model = torch.load(default_models["denoise"])
+    denoise_model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True)
+except: print("Loading model failed")
+denoise_model.eval()
+
+sr_model = SwinIR(upscale=4, in_chans=3, img_size=64, window_size=8,
+            img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
+            mlp_ratio=2, upsampler='nearest+conv', resi_connection='1conv').to(device)
+param_key_g = 'params_ema'
+try:
+    pretrained_model = torch.load(default_models["sr"])
+    sr_model.load_state_dict(pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model, strict=True)
+except: print("Loading model failed")
+sr_model.eval()
+
+
+def sr(input_img):
+
+    window_size = 8
+        # read image
+    img_lq = input_img.astype(np.float32) / 255.
+    img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1))  # HCW-BGR to CHW-RGB
+    img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device)  # CHW-RGB to NCHW-RGB
+
+    # inference
+    with torch.no_grad():
+        # pad input image to be a multiple of window_size
+        _, _, h_old, w_old = img_lq.size()
+        h_pad = (h_old // window_size + 1) * window_size - h_old
+        w_pad = (w_old // window_size + 1) * window_size - w_old
+        img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
+        img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
+        output = sr_model(img_lq)
+        output = output[..., :h_old * 4, :w_old * 4]
+
+    # save image
+    output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+    if output.ndim == 3:
+        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # CHW-RGB to HCW-BGR
+    output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
+
+    return output
+
+def denoise(input_img):
+
+    window_size = 8
+    # read image
+    img_lq = input_img.astype(np.float32) / 255.
+    img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1))  # HCW-BGR to CHW-RGB
+    img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device)  # CHW-RGB to NCHW-RGB
+
+    # inference
+    with torch.no_grad():
+        # pad input image to be a multiple of window_size
+        _, _, h_old, w_old = img_lq.size()
+        h_pad = (h_old // window_size + 1) * window_size - h_old
+        w_pad = (w_old // window_size + 1) * window_size - w_old
+        img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
+        img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
+        output = denoise_model(img_lq)
+        output = output[..., :h_old * 4, :w_old * 4]
+
+    # save image
+    output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+    if output.ndim == 3:
+        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # CHW-RGB to HCW-BGR
+    output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
+    
+    return output
+
+title = " AISeed AI Application Demo "
+description = "# A Demo of Deep Learning for Image Restoration"
+example_list = [["examples/" + example] for example in os.listdir("examples")]
+
+with gr.Blocks() as demo:
+    demo.title = title
+    gr.Markdown(description)
+    with gr.Row():
+        with gr.Column():
+            im = gr.Image(label="Input Image")
+            im_2 = gr.Image(label="Enhanced Image")
+            
+        with gr.Column():
+           
+            btn1 = gr.Button(value="Enhance Resolution")
+            btn1.click(sr, inputs=[im], outputs=[im_2])
+            btn2 = gr.Button(value="Denoise")
+            btn2.click(denoise, inputs=[im], outputs=[im_2])
+            gr.Examples(examples=example_list,
+                inputs=[im],
+                outputs=[im_2])
+    
+if __name__ == "__main__":
+    demo.launch()