Update app.py

app.py

@@ -4,67 +4,87 @@ import os
 from skimage import img_as_ubyte
 import cv2
 import argparse
-
-parser = argparse.ArgumentParser(description='Test Restormer on your own images')
-parser.add_argument('--input_path', default='./temp/image.jpg', type=str, help='Directory of input images or path of single image')
-parser.add_argument('--result_dir', default='./temp/', type=str, help='Directory for restored results')
-parser.add_argument('--task', required=True, type=str, help='Task to run', choices=['Motion_Deblurring',
-                                                                                    'Single_Image_Defocus_Deblurring',
-                                                                                    'Deraining',
-                                                                                    'Real_Denoising',
-                                                                                    'Gaussian_Gray_Denoising',
-                                                                                    'Gaussian_Color_Denoising'])
-
-args = parser.parse_args()
-
-
-task    = args.task
-out_dir = os.path.join(args.result_dir, task)
-
-os.makedirs(out_dir, exist_ok=True)
-
-
-if task == 'Motion_Deblurring':
-    model = torch.jit.load('motion_deblurring.pt')
-elif task == 'Single_Image_Defocus_Deblurring':
-    model = torch.jit.load('single_image_defocus_deblurring.pt')
-elif task == 'Deraining':    
-    model = torch.jit.load('deraining.pt')
-elif task == 'Real_Denoising':
-    model = torch.jit.load('real_denoising.pt')
+import shutil
+import gradio as gr
+from PIL import Image
+
+examples = [['sample1.png'],
+            ['sample2.png']]
+
+inference_on = ['Full Resolution Image', 'Downsampled Image']
+
+title = "Restormer"
+description = """
+Gradio demo for <b>Restormer: Efficient Transformer for High-Resolution Image Restoration</b>, CVPR 2022--ORAL. <a href='https://arxiv.org/abs/2111.09881'>[Paper]</a><a href='https://github.com/swz30/Restormer'>[Github Code]</a>\n 
+<b> Note:</b> Since this demo uses CPU, by default it will run on the downsampled version of the input image (for speedup). But if you want to perform inference on the original input, then choose the option "Full Resolution Image" from the dropdown menu. 
+"""
+##With Restormer, you can perform: (1) Image Denoising, (2) Defocus Deblurring, (3)  Motion Deblurring, and (4) Image Deraining. 
+##To use it, simply upload your own image, or click one of the examples provided below.
+
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.09881'>Restormer: Efficient Transformer for High-Resolution Image Restoration </a> | <a href='https://github.com/swz30/Restormer'>Github Repo</a></p>"
+
+
+def inference(img, task, run_on):
+    if not os.path.exists('temp'):
+      os.system('mkdir temp')
+      
+    if run_on == 'Full Resolution Image':
+      img = img
+    else: # 'Downsampled Image'
+    ####  Resize the longer edge of the input image
+      max_res = 512
+      width, height = img.size
+      if max(width,height) > max_res:
+        scale = max_res /max(width,height)
+        width = int(scale*width)
+        height = int(scale*height)
+        img = img.resize((width,height), Image.ANTIALIAS)
+      
+
+    model = torch.jit.load('deraining.pt')    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
+    model = model.to(device)
+    model.eval()
     
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-# device = torch.device('cpu')
-# stx()
-
-model = model.to(device)
-model.eval()
-
-img_multiple_of = 8
-
-with torch.inference_mode():
-    if torch.cuda.is_available():
-        torch.cuda.ipc_collect()
-        torch.cuda.empty_cache()
+    img_multiple_of = 8
     
-    img = cv2.cvtColor(cv2.imread(args.input_path), cv2.COLOR_BGR2RGB)
-
-    input_ = torch.from_numpy(img).float().div(255.).permute(2,0,1).unsqueeze(0).to(device)
-
-    # Pad the input if not_multiple_of 8
-    h,w = input_.shape[2], input_.shape[3]
-    H,W = ((h+img_multiple_of)//img_multiple_of)*img_multiple_of, ((w+img_multiple_of)//img_multiple_of)*img_multiple_of
-    padh = H-h if h%img_multiple_of!=0 else 0
-    padw = W-w if w%img_multiple_of!=0 else 0
-    input_ = F.pad(input_, (0,padw,0,padh), 'reflect')
-
-    # print(h,w)
-    restored = torch.clamp(model(input_),0,1)
-
-    # Unpad the output
-    restored = img_as_ubyte(restored[:,:,:h,:w].permute(0, 2, 3, 1).cpu().detach().numpy()[0])
-
-    out_path = os.path.join(out_dir, os.path.split(args.input_path)[-1])
-    cv2.imwrite(out_path,cv2.cvtColor(restored, cv2.COLOR_RGB2BGR))
-
-    # print(f"\nRestored images are saved at {out_dir}")
+    with torch.inference_mode():
+        if torch.cuda.is_available():
+            torch.cuda.ipc_collect()
+            torch.cuda.empty_cache()
+        
+        img = cv2.cvtColor(cv2.imread(args.input_path), cv2.COLOR_BGR2RGB)
+    
+        input_ = torch.from_numpy(img).float().div(255.).permute(2,0,1).unsqueeze(0).to(device)
+    
+        # Pad the input if not_multiple_of 8
+        h,w = input_.shape[2], input_.shape[3]
+        H,W = ((h+img_multiple_of)//img_multiple_of)*img_multiple_of, ((w+img_multiple_of)//img_multiple_of)*img_multiple_of
+        padh = H-h if h%img_multiple_of!=0 else 0
+        padw = W-w if w%img_multiple_of!=0 else 0
+        input_ = F.pad(input_, (0,padw,0,padh), 'reflect')
+    
+        restored = torch.clamp(model(input_),0,1)
+    
+        # Unpad the output
+        restored = img_as_ubyte(restored[:,:,:h,:w].permute(0, 2, 3, 1).cpu().detach().numpy()[0])
+        #convert to pil when returning
+  
+    return Image.fromarray(cv2.cvtColor(restored, cv2.COLOR_RGB2BGR))
+    
+gr.Interface(
+    inference,
+    [
+        gr.inputs.Image(type="pil", label="Input"),
+        gr.inputs.Radio(["Deraining"], default="Denoising", label='task'),
+        gr.inputs.Dropdown(choices=inference_on, type="value", default='Downsampled Image', label='Inference on')
+
+    ],
+    gr.outputs.Image(type="pil", label="cleaned and restored"),
+    title=title,
+    description=description,
+    article=article,
+    theme ="huggingface",
+    examples=examples,
+    allow_flagging=False,
+    ).launch(debug=False,enable_queue=True)