Update app.py

app.py

@@ -3,7 +3,7 @@ import sahi
 import torch
 from ultralyticsplus import YOLO, render_model_output
 
-# Download sample images
+# Download images for the demo
 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",
 )
 
-# List of YOLOv8 segmentation models
+# Define available YOLOv8 segmentation models
 model_names = [
     "yolov8n-seg.pt",
     "yolov8s-seg.pt",
@@ -26,18 +26,19 @@ model_names = [
     "yolov8x-seg.pt",
 ]
 
+# Load the initial YOLOv8 model
 current_model_name = "yolov8m-seg.pt"
 model = YOLO(current_model_name)
 
 def yolov8_inference(
     image: gr.Image = None,
-    model_name: str = None,
-    image_size: int = 640,
-    conf_threshold: float = 0.25,
-    iou_threshold: float = 0.45,
+    model_name: gr.Dropdown = None,
+    image_size: gr.Slider = 640,
+    conf_threshold: gr.Slider = 0.25,
+    iou_threshold: gr.Slider = 0.45,
 ):
     """
-    YOLOv8 inference function to return masks and label names for each detected object
+    YOLOv8 inference function
     Args:
         image: Input image
         model_name: Name of the model
@@ -45,85 +46,66 @@ def yolov8_inference(
         conf_threshold: Confidence threshold
         iou_threshold: IOU threshold
     Returns:
-        Object masks, coordinates, and label names
+        Rendered image and mask coordinates with labels
     """
     global model
     global current_model_name
-    
-    # Check if a new model is selected
+    # Switch model if a different one is selected
     if model_name != current_model_name:
         model = YOLO(model_name)
         current_model_name = model_name
     
-    # Set the confidence and IOU thresholds
+    # Set model confidence and IOU thresholds
     model.overrides["conf"] = conf_threshold
     model.overrides["iou"] = iou_threshold
     
-    # Perform model prediction
+    # Perform inference with the YOLO model
     results = model.predict(image, imgsz=image_size, return_outputs=True)
-
-    # Initialize an empty list to store the output
-    output = []
     
-    # Iterate over the results
+    masks = []
     for result in results:
-        # 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})
+        masks.append([result.masks, result.labels])
     
-    return output
+    renders = []
+    for image_results in results:
+        render = render_model_output(
+            model=model, image=image, model_output=image_results
+        )
+        renders.append(render)
 
-# Define Gradio interface inputs and outputs
+    # Return mask coordinates and labels
+    return masks
+
+# Gradio app 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"),
 ]
 
-# 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"
+outputs = gr.Textbox(label="Mask Coordinates and Labels")
 
-# Example images for the interface
+# Example inputs for the Gradio app
 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],
 ]
 
-# Build the Gradio demo app with POST functionality
+# Create the Gradio app interface
 demo_app = gr.Interface(
     fn=yolov8_inference,
     inputs=inputs,
     outputs=outputs,
-    title=title,
+    title="Ultralytics YOLOv8 Segmentation Demo",
     examples=examples,
-    cache_examples=False,  # Set to False to avoid caching issues
-    theme="default",
+    cache_examples=True,
 )
 
-# Launch the app with API-enabled functionality
-demo_app.queue().launch(
-    enable_queue=True,  # Allow for API-style interactions
+# 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