Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import sahi
 import torch
-from ultralyticsplus import YOLO
+from ultralyticsplus import YOLO, render_model_output
 
-# Download example images
+# 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 the model names
+# List of YOLOv8 segmentation models
 model_names = [
     "yolov8n-seg.pt",
     "yolov8s-seg.pt",
@@ -26,7 +26,6 @@ model_names = [
     "yolov8x-seg.pt",
 ]
 
-# Set the initial model
 current_model_name = "yolov8m-seg.pt"
 model = YOLO(current_model_name)
 
@@ -62,18 +61,29 @@ def yolov8_inference(
     
     # Perform model prediction
     results = model.predict(image, imgsz=image_size, return_outputs=True)
-    
-    # Extract masks, coordinates, and label names
+
+    # Initialize an empty list to store the output
     output = []
+    
+    # Iterate over the results
     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})
+        # 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})
     
     return output
 
-# Define Gradio input and output components
+# Define Gradio interface inputs and outputs
 inputs = [
     gr.Image(type="filepath", label="Input Image"),
     gr.Dropdown(
@@ -88,25 +98,26 @@ inputs = [
     gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
 ]
 
-# Output the object masks and label names as a JSON-like format
-outputs = gr.JSON(label="Detected Objects with Masks and Labels")
+# Output is a dictionary containing label names and coordinates of masks or boxes
+outputs = gr.JSON(label="Output Masks and Labels")
 
-# Title and example inputs
 title = "Ultralytics YOLOv8 Segmentation Demo"
+
+# 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 interface
+# Build the Gradio demo app
 demo_app = gr.Interface(
     fn=yolov8_inference,
     inputs=inputs,
     outputs=outputs,
     title=title,
     examples=examples,
-    cache_examples=True,
+    cache_examples=False,  # Set to False to avoid caching issues
     theme="default",
 )