Update app.py

app.py

@@ -3,7 +3,7 @@ import sahi
 import torch
 from ultralyticsplus import YOLO, render_model_output
 
-# Download images for the demo
+# Download sample 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 available YOLOv8 segmentation models
+# List of YOLOv8 segmentation models
 model_names = [
     "yolov8n-seg.pt",
     "yolov8s-seg.pt",
@@ -26,7 +26,6 @@ model_names = [
     "yolov8x-seg.pt",
 ]
 
-# Load the initial YOLOv8 model
 current_model_name = "yolov8m-seg.pt"
 model = YOLO(current_model_name)
 
@@ -38,7 +37,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,68 +45,81 @@ def yolov8_inference(
         conf_threshold: Confidence threshold
         iou_threshold: IOU threshold
     Returns:
-        Rendered image and mask coordinates with labels
+        Object masks, coordinates, and label names
     """
     global model
     global current_model_name
-    # Switch model if a different one is selected
+    
+    # Check if a new model is selected
     if model_name != current_model_name:
         model = YOLO(model_name)
         current_model_name = model_name
     
-    # Set model confidence and IOU thresholds
+    # Set the confidence and IOU thresholds
     model.overrides["conf"] = conf_threshold
     model.overrides["iou"] = iou_threshold
     
-    # Perform inference with the YOLO model
+    # Perform model prediction
     results = model.predict(image, imgsz=image_size, return_outputs=True)
+
+    # Initialize an empty list to store the output
+    output = []
     
-    masks = []
+    # Iterate over the results
     for result in results:
-        masks.append([result.masks, result.labels])
+        # Check if segmentation masks are available
+        if 'masks' in result and result['masks'] is not None:
+            masks = result['masks']['data']
+            for i, (mask, box) in enumerate(zip(masks, result['boxes'])):
+                label = model.names[int(result['boxes']['cls'][i])]
+                mask_coords = mask.tolist()  # Convert mask coordinates to list format
+                output.append({"label": label, "mask_coords": mask_coords})
+        else:
+            # If masks are not available, just extract bounding box information
+            for i, box in enumerate(result['boxes']):
+                label = model.names[int(result['boxes']['cls'][i])]
+                bbox = box['xyxy'].tolist()  # Bounding box coordinates
+                output.append({"label": label, "bbox_coords": bbox})
     
-    renders = []
-    for image_results in results:
-        render = render_model_output(
-            model=model, image=image, model_output=image_results
-        )
-        renders.append(render)
+    return output
 
-    # Return mask coordinates and labels
-    return masks
-
-# Gradio app inputs and outputs
+# Define Gradio interface inputs and outputs
 inputs = [
     gr.Image(type="filepath", label="Input Image"),
-    gr.Dropdown(model_names, value=current_model_name, label="Model type"),
+    gr.Dropdown(
+        model_names,
+        value=current_model_name,
+        label="Model type",
+    ),
     gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
-    gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
+    gr.Slider(
+        minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"
+    ),
     gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
 ]
 
-outputs = gr.Textbox(label="Mask Coordinates and Labels")
+# Output is a dictionary containing label names and coordinates of masks or boxes
+outputs = gr.JSON(label="Output Masks and Labels")
+
+title = "Ultralytics YOLOv8 Segmentation Demo"
 
-# Example inputs for the Gradio app
+# Example images for the interface
 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 app interface
+# Build the Gradio demo app
 demo_app = gr.Interface(
     fn=yolov8_inference,
     inputs=inputs,
     outputs=outputs,
-    title="Ultralytics YOLOv8 Segmentation Demo",
+    title=title,
     examples=examples,
-    cache_examples=True,
+    cache_examples=False,  # Set to False to avoid caching issues
+    theme="default",
 )
 
-# Launch the Gradio app
-demo_app.launch(
-    debug=True,         # Show detailed errors in case of issues
-    server_name="0.0.0.0",  # Host on all IPs
-    server_port=7860,       # Custom port for accessing the app
-    share=True              # To make the app accessible from a URL
-)
+# Launch the app
+demo_app.queue().launch(debug=True)