Update app.py

app.py

@@ -16,6 +16,10 @@ def load_model(model_name):
         model_path = "yolov9c.pt"
     elif model_name == "yolov9e":
         model_path = "yolov9e.pt"
+    elif model_name == "yolov8n":
+        model_path = "yolov8n.pt"
+    elif model_name == "yolov8n-seg":
+        model_path = "yolov8n-seg.pt"
     else:
         raise ValueError(f"Invalid model name: {model_name}")
     
@@ -26,12 +30,12 @@ def predict_image(img, conf_threshold, iou_threshold, task="detection", model_na
     if task == "segmentation":
         if not model_name:
             model_name = "yolov9c-seg"
-        elif model_name not in ["yolov9c-seg", "yolov9e-seg"]:
+        elif model_name not in ["yolov9c-seg", "yolov9e-seg", "yolov8n-seg"]:
             raise ValueError(f"Invalid model name for segmentation: {model_name}")
     elif task == "detection":
         if not model_name:
             model_name = "yolov9c"
-        elif model_name not in ["yolov9c", "yolov9e"]:
+        elif model_name not in ["yolov9c", "yolov9e", "yolov8n"]:
             raise ValueError(f"Invalid model name for detection: {model_name}")
     else:
         raise ValueError(f"Invalid task: {task}. Choose either 'segmentation' or 'detection'.")
@@ -63,13 +67,14 @@ image_iface = gr.Interface(
         gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
         gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
         gr.Dropdown(choices=["detection", "segmentation"], value="detection", label="Task"),
-        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov9c-seg", "yolov9e-seg"], value="yolov9c", label="Model"),
+        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov8n", "yolov9c-seg", "yolov9e-seg", "yolov8n-seg"], value="yolov9c", label="Model"),
     ],
     outputs=gr.Image(type="pil", label="Result"),
     title="X509",
     description="Upload images for inference. Choose task and corresponding model.",
     examples=[
         ["cars.jpg", 0.25, 0.45, "detection", "yolov9c"],
+        ["cars.jpg", 0.25, 0.45, "segmentation", "yolov9c-seg"],
     ],
 )
 
@@ -78,12 +83,12 @@ def predict_video(video_path, conf_threshold, iou_threshold, task="detection", m
     if task == "segmentation":
         if not model_name:
             model_name = "yolov9c-seg"
-        elif model_name not in ["yolov9c-seg", "yolov9e-seg"]:
+        elif model_name not in ["yolov9c-seg", "yolov9e-seg", "yolov8n-seg"]:
             raise ValueError(f"Invalid model name for segmentation: {model_name}")
     elif task == "detection":
         if not model_name:
             model_name = "yolov9c"
-        elif model_name not in ["yolov9c", "yolov9e"]:
+        elif model_name not in ["yolov9c", "yolov9e", "yolov8n"]:
             raise ValueError(f"Invalid model name for detection: {model_name}")
     else:
         raise ValueError(f"Invalid task: {task}. Choose either 'segmentation' or 'detection'.")
@@ -146,13 +151,14 @@ video_iface = gr.Interface(
         gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
         gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
         gr.Dropdown(choices=["detection", "segmentation"], value="detection", label="Task"),
-        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov9c-seg", "yolov9e-seg"], value="yolov9c", label="Model"),
+        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov8n", "yolov9c-seg", "yolov9e-seg", "yolov8n-seg"], value="yolov9c", label="Model"),
     ],
     outputs=gr.File(label="Result"),
     title="X509",
     description="Upload video for inference. Choose task and corresponding model.",
     examples=[
-        ["VID_20240517112011.mp4", 0.25, 0.45, "detection", "yolov9c"],
+        ["VID_20240517112011.mp4", 0.25, 0.45, "detection", "yolov8n"],
+        ["VID_20240517112011.mp4", 0.25, 0.45, "segmentation", "yolov8n-seg"],
     ]
 )