Update app.py

app.py

@@ -106,7 +106,15 @@ def patch2img(outs, idxes, sr_size, scale=4, crop_size=512):
     return (preds / count_mt).to(outs.device)
 
 
-os.makedirs('./results', exist_ok=True)
+def load_img (filename, norm=True,):
+    img = np.array(Image.open(filename).convert("RGB"))
+    h, w = img.shape[:2]
+    
+    if norm:
+        img = img.astype(np.float32) / 255.
+        
+    return img
+
 
 def inference(image, upscale, large_input_flag, color_fix):
 	upscale = int(upscale) # convert type to int
@@ -117,17 +125,16 @@ def inference(image, upscale, large_input_flag, color_fix):
 
 	model = set_safmn(upscale)
 
-	img = cv2.imread(str(image), cv2.IMREAD_COLOR)
-	print(f'input size: {img.shape}')
+    img = np.array(image)
+    img = img / 255.
+    img = img.astype(np.float32)
 
-	# img2tensor
-	img = img.astype(np.float32) / 255.
-	img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
-	img = img.unsqueeze(0).to(device)
+    # img2tensor
+    y = torch.tensor(img).permute(2,0,1).unsqueeze(0).to(device)
 
 	# inference
 	if large_input_flag:
-		patches, idx, size = img2patch(img, scale=upscale)
+		patches, idx, size = img2patch(y, scale=upscale)
 		with torch.no_grad():
 			n = len(patches)
 			outs = []
@@ -147,24 +154,19 @@ def inference(image, upscale, large_input_flag, color_fix):
 		output = patch2img(output, idx, size, scale=upscale)
 	else:
 		with torch.no_grad():
-			output = model(img)
+			output = model(y)
 
 	# color fix
 	if color_fix:
-		img = F.interpolate(img, scale_factor=upscale, mode='bilinear')
-		output = wavelet_reconstruction(output, img)
+		y = F.interpolate(y, scale_factor=upscale, mode='bilinear')
+		output = wavelet_reconstruction(output, y)
 	# tensor2img
 	output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
 	if output.ndim == 3:
 		output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
 	output = (output * 255.0).round().astype(np.uint8)
 
-	# save restored img
-	save_path = f'results/out.png'
-	cv2.imwrite(save_path, output)
-
-	output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-	return output, save_path
+	return (image, Image.fromarray(output))
 
 
 
@@ -194,21 +196,16 @@ article = "<p style='text-align: center'><a href='https://eduardzamfir.github.io
 
 #### Image,Prompts examples
 examples = [
-            ['images/0801x4.png'],
-            ['images/0840x4.png'],
-            ['images/0841x4.png'],
-            ['images/0870x4.png'],
-            ['images/0878x4.png'],
-            ['images/0884x4.png'],
-            ['images/0900x4.png'],
-            ['images/img002x4.png'],
-            ['images/img003x4.png'],
-            ['images/img004x4.png'],
-            ['images/img035x4.png'],
-            ['images/img053x4.png'],
-            ['images/img064x4.png'],
-            ['images/img083x4.png'],
-            ['images/img092x4.png'],
+            ['real_testdata/004.png'],
+            ['real_testdata/005.png'],
+            ['real_testdata/010.png'],
+            ['real_testdata/015.png'],
+            ['real_testdata/025.png'],
+            ['real_testdata/030.png'],
+            ['real_testdata/034.png'],
+            ['real_testdata/044.png'],
+            ['real_testdata/041.png'],
+            ['real_testdata/054.png'],
             ]
 
 css = """
@@ -220,7 +217,7 @@ css = """
 """
 
 demo = gr.Interface(
-    fn=process_img,
+    fn=inference,
     inputs=[
         gr.Image(type="pil", label="Input", value="real_testdata/004.png"),
         gr.Number(default=2, label="Upscaling factor (up to 4)"),