Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import sahi
 import torch
-from ultralyticsplus import YOLO, render_model_output
+from ultralyticsplus import YOLO
 
-# Images
+# Download example images
 sahi.utils.file.download_from_url(
     "https://raw.githubusercontent.com/kadirnar/dethub/main/data/images/highway.jpg",
     "highway.jpg",
@@ -17,7 +17,7 @@ sahi.utils.file.download_from_url(
     "zidane.jpg",
 )
 
-
+# Define the model names
 model_names = [
     "yolov8n-seg.pt",
     "yolov8s-seg.pt",
@@ -26,10 +26,10 @@ model_names = [
     "yolov8x-seg.pt",
 ]
 
+# Set the initial model
 current_model_name = "yolov8m-seg.pt"
 model = YOLO(current_model_name)
 
-
 def yolov8_inference(
     image: gr.Image = None,
     model_name: gr.Dropdown = None,
@@ -38,7 +38,7 @@ def yolov8_inference(
     iou_threshold: gr.Slider = 0.45,
 ):
     """
-    YOLOv8 inference function
+    YOLOv8 inference function to return masks and label names for each detected object
     Args:
         image: Input image
         model_name: Name of the model
@@ -46,29 +46,34 @@ def yolov8_inference(
         conf_threshold: Confidence threshold
         iou_threshold: IOU threshold
     Returns:
-        Rendered image
+        Object masks, coordinates, and label names
     """
     global model
     global current_model_name
+    
+    # Check if a new model is selected
     if model_name != current_model_name:
         model = YOLO(model_name)
         current_model_name = model_name
+    
+    # Set the confidence and IOU thresholds
     model.overrides["conf"] = conf_threshold
     model.overrides["iou"] = iou_threshold
+    
+    # Perform model prediction
     results = model.predict(image, imgsz=image_size, return_outputs=True)
-    masks=[]
-    for i in results:
-        masks.append([i.masks,i.labels])
-    renders = []
-    for image_results in model.predict(image, imgsz=image_size, return_outputs=True):
-        render = render_model_output(
-            model=model, image=image, model_output=image_results
-        )
-        renders.append(render)
-
-    return masks
-
+    
+    # Extract masks, coordinates, and label names
+    output = []
+    for result in results:
+        for mask, box in zip(result.masks.xy, result.boxes):
+            label = model.names[int(box.cls[0])]
+            mask_coords = mask.tolist()  # Convert mask coordinates to list format
+            output.append({"label": label, "mask_coords": mask_coords})
+    
+    return output
 
+# Define Gradio input and output components
 inputs = [
     gr.Image(type="filepath", label="Input Image"),
     gr.Dropdown(
@@ -83,14 +88,18 @@ inputs = [
     gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
 ]
 
-outputs = gr.Image(type="filepath", label="Output Image")
-title = "Ultralytics YOLOv8 Segmentation Demo"
+# Output the object masks and label names as a JSON-like format
+outputs = gr.JSON(label="Detected Objects with Masks and Labels")
 
+# Title and example inputs
+title = "Ultralytics YOLOv8 Segmentation Demo"
 examples = [
     ["zidane.jpg", "yolov8m-seg.pt", 640, 0.6, 0.45],
     ["highway.jpg", "yolov8m-seg.pt", 640, 0.25, 0.45],
     ["small-vehicles1.jpeg", "yolov8m-seg.pt", 640, 0.25, 0.45],
 ]
+
+# Create the Gradio interface
 demo_app = gr.Interface(
     fn=yolov8_inference,
     inputs=inputs,
@@ -100,4 +109,6 @@ demo_app = gr.Interface(
     cache_examples=True,
     theme="default",
 )
-demo_app.queue().launch(debug=True)
+
+# Launch the app
+demo_app.queue().launch(debug=True)