Update app.py

app.py

@@ -4,19 +4,21 @@ import cv2
 import numpy as np
 import os
 import tempfile
+import base64
+from io import BytesIO
 from PIL import Image
 
-def perform(image):
+def perform(image, task: str="detect"):
     # Predict with the model
     results = model(image, save=True)  # predict on an image
     
     # get the saved prediction
     prediction_image = Image.open(os.path.join(results[0].save_dir, os.listdir(results[0].save_dir)[0]))
     
-    # Save the prediction image to a temporary file for public URL
-    temp_dir = tempfile.mkdtemp()
-    temp_path = os.path.join(temp_dir, "prediction.jpg")
-    prediction_image.save(temp_path)
+    # Convert PIL image to base64
+    buffered = BytesIO()
+    prediction_image.save(buffered, format="JPEG")
+    img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
     
     # initialize the output list to save the bbox
     outputs = []
@@ -28,8 +30,11 @@ def perform(image):
         confs = float(result.boxes.conf)  # confidence score of each box
         outputs.append({"xywh": xywh, "conf": confs, "name": name})
     
-    # return the image path (for public URL) and the bbox
-    return temp_path, outputs
+    # Add base64 image to outputs
+    outputs.append({"image_base64": f"data:image/jpeg;base64,{img_base64}"})
+    
+    # return the PIL image (for display) and the outputs with base64
+    return prediction_image, outputs
 
 def load_model(model_name: str = "yolo11n.pt"):
     """Model loading"""
@@ -47,7 +52,7 @@ model = load_model()
 demo = gr.Interface(
     fn=perform,
     inputs=[gr.Image(type="pil")],  # PIL Image format
-    outputs=[gr.Image(type="filepath"), gr.JSON()],  # Changed to filepath for public URL
+    outputs=[gr.Image(type="pil"), gr.JSON()],  # Back to PIL - Gradio handles public URL automatically
     title="Object Detection using YOLO",
     description="Upload an image and get image with bbox in it",
     examples=None  # You can add example images here if needed
@@ -55,4 +60,4 @@ demo = gr.Interface(
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()  # Added share=True for public access
+    demo.launch()