Updated Code

app.py

@@ -6,42 +6,11 @@ import supervision as sv
 from inference import get_model
 
 MARKDOWN = """
-<h1 style='text-align: center'>YOLO-ARENA 🏟️</h1>
-Welcome to YOLO-Arena! This demo showcases the performance of various YOLO models 
-pre-trained on the COCO dataset.
-- **YOLOv8**
-  <div style="display: flex; align-items: center;">
-    <a href="https://github.com/ultralytics/ultralytics" style="margin-right: 10px;">
-      <img src="https://badges.aleen42.com/src/github.svg">
-    </a>
-    <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov8-object-detection-on-custom-dataset.ipynb" style="margin-right: 10px;">
-      <img src="https://colab.research.google.com/assets/colab-badge.svg">
-    </a>
-  </div>
-- **YOLOv9**
-  <div style="display: flex; align-items: center;">
-    <a href="https://github.com/WongKinYiu/yolov9" style="margin-right: 10px;">
-      <img src="https://badges.aleen42.com/src/github.svg">
-    </a>
-    <a href="https://arxiv.org/abs/2402.13616" style="margin-right: 10px;">
-      <img src="https://img.shields.io/badge/arXiv-2402.13616-b31b1b.svg">
-    </a>
-    <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov9-object-detection-on-custom-dataset.ipynb" style="margin-right: 10px;">
-      <img src="https://colab.research.google.com/assets/colab-badge.svg">
-    </a>
-  </div>
-- **YOLOv10**
-  <div style="display: flex; align-items: center;">
-    <a href="https://github.com/THU-MIG/yolov10" style="margin-right: 10px;">
-      <img src="https://badges.aleen42.com/src/github.svg">
-    </a>
-    <a href="https://arxiv.org/abs/2405.14458" style="margin-right: 10px;">
-      <img src="https://img.shields.io/badge/arXiv-2405.14458-b31b1b.svg">
-    </a>
-    <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov10-object-detection-on-custom-dataset.ipynb" style="margin-right: 10px;">
-      <img src="https://colab.research.google.com/assets/colab-badge.svg">
-    </a>
-  </div>
+<h1 style='text-align: center'>Detect Something 💫</h1>
+Welcome to Segment Something! Your on the go demo for instance segmentation. 🚀
+
+<h2 style='text-align: center'>Matthias Bartolo</h2>
+
 Powered by Roboflow [Inference](https://github.com/roboflow/inference) and 
 [Supervision](https://github.com/roboflow/supervision). 🔥
 """
@@ -52,9 +21,9 @@ IMAGE_EXAMPLES = [
     ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.3, 0.3, 0.1],
 ]
 
-YOLO_V8_MODEL = get_model(model_id="coco/8")
-YOLO_V9_MODEL = get_model(model_id="coco/17")
-YOLO_V10_MODEL = get_model(model_id="coco/22")
+YOLO_V8N_MODEL = get_model(model_id="yolov8n-640")
+YOLO_V8S_MODEL = get_model(model_id="yolov8s-640")
+YOLO_V8M_MODEL = get_model(model_id="yolov8m-640")
 
 LABEL_ANNOTATORS = sv.LabelAnnotator(text_color=sv.Color.black())
 BOUNDING_BOX_ANNOTATORS = sv.BoundingBoxAnnotator()
@@ -102,26 +71,30 @@ def process_image(
     yolo_v10_confidence_threshold: float,
     iou_threshold: float
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    yolo_v8_annotated_image = detect_and_annotate(
-        YOLO_V8_MODEL, input_image, yolo_v8_confidence_threshold, iou_threshold)
-    yolo_v9_annotated_image = detect_and_annotate(
-        YOLO_V9_MODEL, input_image, yolo_v9_confidence_threshold, iou_threshold)
-    yolo_10_annotated_image = detect_and_annotate(
-        YOLO_V10_MODEL, input_image, yolo_v10_confidence_threshold, iou_threshold)
+    # Validate iou_threshold before using it
+    if iou_threshold is None or not isinstance(iou_threshold, float):
+        iou_threshold = 0.3  # Default value, adjust as necessary
+
+    yolo_v8n_annotated_image = detect_and_annotate(
+        YOLO_V8N_MODEL, input_image, yolo_v8_confidence_threshold, iou_threshold)
+    yolo_v8s_annotated_image = detect_and_annotate(
+        YOLO_V8S_MODEL, input_image, yolo_v9_confidence_threshold, iou_threshold)
+    yolo_8m_annotated_image = detect_and_annotate(
+        YOLO_V8M_MODEL, input_image, yolo_v10_confidence_threshold, iou_threshold)
 
     return (
-        yolo_v8_annotated_image,
-        yolo_v9_annotated_image,
-        yolo_10_annotated_image
+        yolo_v8n_annotated_image,
+        yolo_v8s_annotated_image,
+        yolo_8m_annotated_image
     )
 
 
-yolo_v8_confidence_threshold_component = gr.Slider(
+yolo_v8N_confidence_threshold_component = gr.Slider(
     minimum=0,
     maximum=1.0,
     value=0.3,
     step=0.01,
-    label="YOLOv8 Confidence Threshold",
+    label="YOLOv8N Confidence Threshold",
     info=(
         "The confidence threshold for the YOLO model. Lower the threshold to "
         "reduce false negatives, enhancing the model's sensitivity to detect "
@@ -129,12 +102,12 @@ yolo_v8_confidence_threshold_component = gr.Slider(
         "positives, preventing the model from identifying objects it shouldn't."
     ))
 
-yolo_v9_confidence_threshold_component = gr.Slider(
+yolo_v8S_confidence_threshold_component = gr.Slider(
     minimum=0,
     maximum=1.0,
     value=0.3,
     step=0.01,
-    label="YOLOv9 Confidence Threshold",
+    label="YOLOv8S Confidence Threshold",
     info=(
         "The confidence threshold for the YOLO model. Lower the threshold to "
         "reduce false negatives, enhancing the model's sensitivity to detect "
@@ -142,12 +115,12 @@ yolo_v9_confidence_threshold_component = gr.Slider(
         "positives, preventing the model from identifying objects it shouldn't."
     ))
 
-yolo_v10_confidence_threshold_component = gr.Slider(
+yolo_v8M_confidence_threshold_component = gr.Slider(
     minimum=0,
     maximum=1.0,
     value=0.3,
     step=0.01,
-    label="YOLOv10 Confidence Threshold",
+    label="YOLOv8M Confidence Threshold",
     info=(
         "The confidence threshold for the YOLO model. Lower the threshold to "
         "reduce false negatives, enhancing the model's sensitivity to detect "
@@ -174,27 +147,27 @@ with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
     with gr.Accordion("Configuration", open=False):
         with gr.Row():
-            yolo_v8_confidence_threshold_component.render()
-            yolo_v9_confidence_threshold_component.render()
-            yolo_v10_confidence_threshold_component.render()
+            yolo_v8N_confidence_threshold_component.render()
+            yolo_v8S_confidence_threshold_component.render()
+            yolo_v8M_confidence_threshold_component.render()
         iou_threshold_component.render()
     with gr.Row():
         input_image_component = gr.Image(
             type='pil',
             label='Input'
         )
-        yolo_v8_output_image_component = gr.Image(
+        yolo_v8n_output_image_component = gr.Image(
             type='pil',
-            label='YOLOv8'
+            label='YOLOv8N'
         )
     with gr.Row():
-        yolo_v9_output_image_component = gr.Image(
+        yolo_v8s_output_image_component = gr.Image(
             type='pil',
-            label='YOLOv9'
+            label='YOLOv8S'
         )
-        yolo_v10_output_image_component = gr.Image(
+        yolo_v8m_output_image_component = gr.Image(
             type='pil',
-            label='YOLOv10'
+            label='YOLOv8M'
             )
     submit_button_component = gr.Button(
         value='Submit',
@@ -206,15 +179,15 @@ with gr.Blocks() as demo:
         examples=IMAGE_EXAMPLES,
         inputs=[
             input_image_component,
-            yolo_v8_confidence_threshold_component,
-            yolo_v9_confidence_threshold_component,
-            yolo_v10_confidence_threshold_component,
+            yolo_v8N_confidence_threshold_component,
+            yolo_v8S_confidence_threshold_component,
+            yolo_v8M_confidence_threshold_component,
             iou_threshold_component
         ],
         outputs=[
-            yolo_v8_output_image_component,
-            yolo_v9_output_image_component,
-            yolo_v10_output_image_component
+            yolo_v8n_output_image_component,
+            yolo_v8s_output_image_component,
+            yolo_v8m_output_image_component
         ]
     )
 
@@ -222,15 +195,15 @@ with gr.Blocks() as demo:
         fn=process_image,
         inputs=[
             input_image_component,
-            yolo_v8_confidence_threshold_component,
-            yolo_v9_confidence_threshold_component,
-            yolo_v10_confidence_threshold_component,
+            yolo_v8N_confidence_threshold_component,
+            yolo_v8S_confidence_threshold_component,
+            yolo_v8M_confidence_threshold_component,
             iou_threshold_component
         ],
         outputs=[
-            yolo_v8_output_image_component,
-            yolo_v9_output_image_component,
-            yolo_v10_output_image_component
+            yolo_v8n_output_image_component,
+            yolo_v8s_output_image_component,
+            yolo_v8m_output_image_component
         ]
     )