fix jpeg error

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import cv2
 import torch
 import urllib.request
+import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 
@@ -46,11 +47,14 @@ def update(slider, img):
 
     output = prediction.cpu().numpy()
 
+    formatted = (output * 255 / np.max(output)).astype('uint8')
+    out_im = Image.fromarray(formatted)
+
     #out_im = Image.fromarray(output)
-    #out_im.convert('RGB').save("temp/image_depth.jpeg", "JPEG")
+    out_im.save("temp/image_depth.jpeg", "JPEG")
 
     #cv2.imwrite("temp/image_depth.jpeg", output)
-    plt.imsave("temp/image_depth.jpeg", output)
+    #plt.imsave("temp/image_depth.jpeg", output)
 
     return f'temp/image_depth.jpeg'