Update app.py

app.py

@@ -2,25 +2,34 @@ import gradio as gr
 from ultralytics import YOLO
 import cv2
 import numpy as np
-import os 
+import os
+import tempfile
 from PIL import Image
 
-def perform(image):
-    """Perform the task of detection"""
+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)
+    
     # initialize the output list to save the bbox
     outputs = []
+    
     # Access the results
     for result in results[0]:
         xywh = result.boxes.xyxy[0].numpy().tolist()  # top-left-x, top-left-y, bottom-right-x, bottom-right-y
-        name = result.names[result.boxes.cls.item()] # classes
+        name = result.names[result.boxes.cls.item()]  # classes
         confs = float(result.boxes.conf)  # confidence score of each box
-        outputs.append({"xywh":xywh ,"conf":confs, "name":name})
-    # return the image and the bbox
-    return prediction_image, outputs
+        outputs.append({"xywh": xywh, "conf": confs, "name": name})
+    
+    # return the image path (for public URL) and the bbox
+    return temp_path, outputs
 
 def load_model(model_name: str = "yolo11n.pt"):
     """Model loading"""
@@ -29,7 +38,6 @@ def load_model(model_name: str = "yolo11n.pt"):
         model = YOLO(model_name)
     except Exception as e:
         print(f"Error loading the model {e}")
-        return None
     return model
 
 # loading the model
@@ -39,7 +47,7 @@ model = load_model()
 demo = gr.Interface(
     fn=perform,
     inputs=[gr.Image(type="pil")],  # PIL Image format
-    outputs=[gr.Image(type="pil"), gr.JSON()],
+    outputs=[gr.Image(type="filepath"), gr.JSON()],  # Changed to filepath for public URL
     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
@@ -47,5 +55,4 @@ demo = gr.Interface(
 
 # Launch the app
 if __name__ == "__main__":
-
-    demo.launch()
+    demo.launch()  # Added share=True for public access