Update app.py

app.py

@@ -4,6 +4,7 @@ import PIL.Image
 from PIL import Image
 from io import BytesIO
 import requests
+import torch
 import torch.nn as nn
 import os
 import tempfile
@@ -42,7 +43,25 @@ def add_margin(pil_img, top, right, bottom, left, color):
     result.paste(pil_img, (left, top))
     return result
 
+def tensor_to_pil(tensor):
+    """
+    Convert a tensor to a PIL Image.
+    """
+    tensor = tensor.cpu().clamp(0, 1)
+    array = tensor.numpy().transpose(1, 2, 0)
+    return Image.fromarray((array * 255).astype('uint8'))
+
 def inference(image_path_or_url, learn):
+    """
+    Perform inference on an image from a local path or a URL.
+
+    Parameters:
+    image_path_or_url (str): Path to the local image or URL of the image.
+    learn (Learner): The trained model.
+
+    Returns:
+    PIL.Image.Image: The high-resolution image generated by the model.
+    """
     if image_path_or_url.startswith('http://') or image_path_or_url.startswith('https://'):
         response = requests.get(image_path_or_url)
         img = PIL.Image.open(BytesIO(response.content)).convert("RGB")
@@ -53,7 +72,7 @@ def inference(image_path_or_url, learn):
     im_new.save("test.jpg", quality=95)
     img = open_image("test.jpg")
     p, img_hr, b = learn.predict(img)
-    return img_hr.to_pil()
+    return tensor_to_pil(img_hr)
 
 # Streamlit application
 st.title("Image Inference with Fastai")