Update demo_gradio.py

demo_gradio.py

@@ -6,10 +6,8 @@
 import torch
 import torch.nn.functional as F
 import os
-from runpy import run_path
 from skimage import img_as_ubyte
 import cv2
-from tqdm import tqdm
 import argparse
 
 parser = argparse.ArgumentParser(description='Test Restormer on your own images')
@@ -25,45 +23,36 @@ parser.add_argument('--task', required=True, type=str, help='Task to run', choic
 args = parser.parse_args()
 
 
-def get_weights_and_parameters(task, parameters):
-    if task == 'Motion_Deblurring':
-        weights = os.path.join('Motion_Deblurring', 'pretrained_models', 'motion_deblurring.pth')
-    elif task == 'Single_Image_Defocus_Deblurring':
-        weights = os.path.join('Defocus_Deblurring', 'pretrained_models', 'single_image_defocus_deblurring.pth')
-    elif task == 'Deraining':
-        weights = os.path.join('Deraining', 'pretrained_models', 'deraining.pth')
-    elif task == 'Real_Denoising':
-        weights = os.path.join('Denoising', 'pretrained_models', 'real_denoising.pth')
-        parameters['LayerNorm_type'] =  'BiasFree'
-    return weights, parameters
-
 task    = args.task
 out_dir = os.path.join(args.result_dir, task)
 
 os.makedirs(out_dir, exist_ok=True)
 
-# Get model weights and parameters
-parameters = {'inp_channels':3, 'out_channels':3, 'dim':48, 'num_blocks':[4,6,6,8], 'num_refinement_blocks':4, 'heads':[1,2,4,8], 'ffn_expansion_factor':2.66, 'bias':False, 'LayerNorm_type':'WithBias', 'dual_pixel_task':False}
-weights, parameters = get_weights_and_parameters(task, parameters)
-
-load_arch = run_path(os.path.join('basicsr', 'models', 'archs', 'restormer_arch.py'))
-model = load_arch['Restormer'](**parameters)
 
+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')
+    
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 # device = torch.device('cpu')
+# stx()
 
 model = model.to(device)
-checkpoint = torch.load(weights)
-model.load_state_dict(checkpoint['params'])
-
-model.eval()
 
 
 img_multiple_of = 8
 
 
 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)
@@ -74,10 +63,14 @@ with torch.inference_mode():
     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])
 
-    cv2.imwrite(os.path.join(out_dir, os.path.split(args.input_path)[-1]),cv2.cvtColor(restored, cv2.COLOR_RGB2BGR))
+    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}")