commit from ssq

app.py

@@ -13,7 +13,7 @@ model_restoration.eval()
 
 factor = 8
 def predict(input_img):
-    img = np.float32(load_img(file_))/255.
+    img = np.float32(load_img(input_img))/255.
     img = torch.from_numpy(img).permute(2,0,1)
     input_ = img.unsqueeze(0).cuda()
 
@@ -23,13 +23,13 @@ def predict(input_img):
     padh = H-h if h%factor!=0 else 0
     padw = W-w if w%factor!=0 else 0
     input_ = F.pad(input_, (0,padw,0,padh), 'reflect')
-    return input_
-    prediction = model_restoration(input_)
+    
+    restored = model_restoration(input_)
     output_path = "restored.png"
-    result_dir = restored[:,:,:h,:w]
+    restored = restored[:,:,:h,:w]
     restored = torch.clamp(restored,0,1).cpu().detach().permute(0, 2, 3, 1).squeeze(0).numpy()
 
-    save_img(result_dir, img_as_ubyte(restored))
+    save_img(output_path, img_as_ubyte(restored))
 
 example_images = [
     "examples/example.jpeg",