pytorch/YOLOv5

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from PIL import Image
-import os 
+
 # Images
 torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
 torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', 'bus.jpg')
@@ -9,19 +9,18 @@ torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master
 # Model
 model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # force_reload=True to update
 
+
 def yolo(im, size=640):
     g = (size / max(im.size))  # gain
     im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize
 
     results = model(im)  # inference
     results.render()  # updates results.imgs with boxes and labels
-    results.save()
-    os.system("ls")
-    return "out.png"
+    return Image.fromarray(results.imgs[0])
 
 
 inputs = gr.inputs.Image(type='pil', label="Original Image")
-outputs = gr.outputs.Image(type="file", label="Output Image")
+outputs = gr.outputs.Image(type="pil", label="Output Image")
 
 title = "YOLOv5"
 description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use."