working with segemntation

download_models.py

@@ -137,13 +137,21 @@ def download_sam2_models():
 def download_yolo_models():
     """Download default YOLO models to models directory."""
     print("\n--- Setting up YOLO models ---")
+    print("  Downloading both detection and segmentation models...")
     
     try:
         from ultralytics import YOLO
         import torch
         
-        # Default YOLO models to download
-        yolo_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"]
+        # Default YOLO models to download (both detection and segmentation)
+        yolo_models = [
+            "yolov8n.pt",      # Detection models
+            "yolov8s.pt", 
+            "yolov8m.pt",
+            "yolov8n-seg.pt",  # Segmentation models
+            "yolov8s-seg.pt",
+            "yolov8m-seg.pt"
+        ]
         models_dir = Path(__file__).parent / "models" / "yolo"
         
         for model_name in yolo_models:
@@ -205,8 +213,13 @@ def download_yolo_models():
         success = all((models_dir / model).exists() for model in yolo_models)
         if success:
             print("✓ YOLO models setup complete!")
+            print("  Available detection models: yolov8n.pt, yolov8s.pt, yolov8m.pt")
+            print("  Available segmentation models: yolov8n-seg.pt, yolov8s-seg.pt, yolov8m-seg.pt")
         else:
-            print("⚠ Some YOLO models may be missing")
+            missing_models = [model for model in yolo_models if not (models_dir / model).exists()]
+            print("⚠ Some YOLO models may be missing:")
+            for model in missing_models:
+                print(f"    - {model}")
         return success
         
     except ImportError:
@@ -234,6 +247,12 @@ def update_config_file():
         updated_content = content.replace(
             'yolo_model: "yolov8n.pt"',
             'yolo_model: "models/yolo/yolov8n.pt"'
+        ).replace(
+            'yolo_detection_model: "models/yolo/yolov8n.pt"',
+            'yolo_detection_model: "models/yolo/yolov8n.pt"'
+        ).replace(
+            'yolo_segmentation_model: "models/yolo/yolov8n-seg.pt"',
+            'yolo_segmentation_model: "models/yolo/yolov8n-seg.pt"'
         ).replace(
             'sam2_checkpoint: "../checkpoints/sam2.1_hiera_large.pt"',
             'sam2_checkpoint: "models/sam2/checkpoints/sam2.1_hiera_large.pt"'