Update app.py

app.py

@@ -56,8 +56,11 @@ app = Flask(__name__)
 
 @app.route('/')
 def index():
+
     '''return '<iframe frameBorder="0" height="100%" src="{}/?__dark-theme={}" width="100%"></iframe>'.format(
     os.getenv('INACCEL_URL'),request.args.get('__dark-theme', 'false'))'''
+    model = torch.hub.load('ultralytics/yolov5', 'custom', path='best2.pt', force_reload=True)  # force_reload to recache
+    
     if request.files.get("image"):
         image_file = request.files["image"]
         image_bytes = image_file.read()
@@ -65,9 +68,14 @@ def index():
         img = Image.open(io.BytesIO(image_bytes))
 
         results = model(img, size=640)  # reduce size=320 for faster inference
-        return results.pandas().xyxy[0].to_json(orient="records")
+        results.imgs # array of original images (as np array) passed to model for inference
+        results.render()  # updates results.imgs with boxes and labels
+        for img in results.imgs:
+            buffered = BytesIO()
+            img_base64 = Image.fromarray(img)
+            img_base64.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')  # base64 encoded image with results
 
 
 if __name__ == '__main__':
-    model = torch.hub.load('ultralytics/yolov5', 'custom', path='best2.pt', force_reload=True)  # force_reload to recache
     app.run(host='0.0.0.0', port=7860)