Add fix for output to PIL image.

app.py

@@ -1,8 +1,7 @@
+import cv2
 import torch
-from torchvision import transforms
 import gradio as gr
-from random import randint
-from pathlib import Path
+from torchvision import transforms
 from super_image import ImageLoader, EdsrModel, MsrnModel, MdsrModel, AwsrnModel, A2nModel, CarnModel, PanModel, \
     HanModel, DrlnModel, RcanModel
 
@@ -40,17 +39,14 @@ def get_model(model_name, scale):
 
 
 def inference(img, scale_str, model_name):
-    _id = randint(1, 1000)
-    output_dir = Path('./tmp/')
-    output_dir.mkdir(parents=True, exist_ok=True)
-    # output_file = output_dir / ('output_image' + str(_id) + '.jpg')
     scale = int(scale_str.replace('x', ''))
     model = get_model(model_name, scale)
     inputs = ImageLoader.load_image(img)
     preds = model(inputs)
-    # output_file_str = str(output_file.resolve())
-    # ImageLoader.save_image(preds, output_file_str)
-    return transforms.ToPILImage(mode='RGB')(preds)
+    preds = preds.data.cpu().numpy()
+    pred = preds[0].transpose((1, 2, 0)) * 255.0
+    pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
+    return transforms.ToPILImage(mode='RGB')(pred)
 
 
 torch.hub.download_url_to_file('http://people.rennes.inria.fr/Aline.Roumy/results/images_SR_BMVC12/input_groundtruth/baby_mini_d3_gaussian.bmp',
@@ -79,5 +75,6 @@ gr.Interface(
         ['baby.bmp', 'x2', 'EDSR-base'],
         ['woman.bmp', 'x3', 'DRLN']
     ],
-    enable_queue=True
+    enable_queue=True,
+    allow_flagging=False,
     ).launch(debug=False)