Update app.py

app.py

@@ -116,23 +116,24 @@ def load_img(filename, norm=True):
 
 
 def inference(image, upscale, large_input_flag, color_fix):
-    if upscale is None or not isinstance(upscale, (int, float)):
-        upscale = 2 
-		
-    upscale = int(upscale)
-    if 0 < upscale < 3:
+	if upscale is None or not isinstance(upscale, (int, float)) or upscale == 3:
 		upscale = 2
+		
+	upscale = int(upscale) # convert type to int
 
 	model = set_safmn(upscale)
 
-	img = np.array(image)
+	img = cv2.imread(str(image), cv2.IMREAD_COLOR)
+	print(f'input size: {img.shape}')
+
+	# img2tensor
 	img = img.astype(np.float32) / 255.
-	y = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
-	y = y.unsqueeze(0).to(device)
+	img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+	img = img.unsqueeze(0).to(device)
 
 	# inference
 	if large_input_flag:
-		patches, idx, size = img2patch(y, scale=upscale)
+		patches, idx, size = img2patch(img, scale=upscale)
 		with torch.no_grad():
 			n = len(patches)
 			outs = []
@@ -152,19 +153,24 @@ def inference(image, upscale, large_input_flag, color_fix):
 		output = patch2img(output, idx, size, scale=upscale)
 	else:
 		with torch.no_grad():
-			output = model(y)
+			output = model(img)
 
 	# color fix
 	if color_fix:
-		y = F.interpolate(img, scale_factor=upscale, mode='bilinear')
-		output = wavelet_reconstruction(output, y)
+		img = F.interpolate(img, scale_factor=upscale, mode='bilinear')
+		output = wavelet_reconstruction(output, img)
 	# 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)
 
-	return (image, output)
+	# 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