Changed from yolov8n-world.pt to yolov8s-world.pt, add error handling

app.py

@@ -5,14 +5,50 @@ import torch
 from ultralytics import SAM, YOLOWorld
 import os
 
-# Initialize models
-sam_model = SAM("mobile_sam.pt")  # Switch to MobileSAM for faster CPU inference
-yolo_model = YOLOWorld("yolov8n-world.pt")  # Nano model for faster detection
+# Initialize models with proper error handling and auto-download
+def initialize_models():
+    """Initialize models with proper error handling."""
+    try:
+        sam_model = SAM("mobile_sam.pt")  # This auto-downloads
+        print("✅ SAM model loaded successfully")
+    except Exception as e:
+        print(f"❌ Error loading SAM model: {e}")
+        raise
+    
+    try:
+        # Try different YOLO-World model names that auto-download
+        yolo_model = YOLOWorld("yolov8s-world.pt")  # Small world model (auto-downloads)
+        print("✅ YOLO-World model loaded successfully")
+        return sam_model, yolo_model
+    except Exception as e:
+        print(f"❌ Error loading YOLO-World model: {e}")
+        try:
+            # Fallback to regular YOLO if YOLO-World fails
+            from ultralytics import YOLO
+            yolo_model = YOLO("yolov8n.pt")  # Regular YOLO nano model
+            print("⚠️ Using regular YOLO model as fallback")
+            return sam_model, yolo_model
+        except Exception as e2:
+            print(f"❌ Fallback YOLO model also failed: {e2}")
+            raise
+
+sam_model, yolo_model = initialize_models()
 
 def detect_motorcycles(first_frame, prompt="motorcycle"):
     """Detect motorcycles in the first frame using YOLO-World and return bounding boxes."""
-    yolo_model.set_classes([prompt])
-    results = yolo_model.predict(first_frame, device="cpu", max_det=2, imgsz=320)  # Force YOLO to use 320 resolution
+    try:
+        # Check if it's YOLO-World model
+        if hasattr(yolo_model, 'set_classes'):
+            yolo_model.set_classes([prompt])
+            results = yolo_model.predict(first_frame, device="cpu", max_det=2, imgsz=320, verbose=False)
+        else:
+            # Regular YOLO model - can't set custom classes, will detect all objects
+            results = yolo_model.predict(first_frame, device="cpu", max_det=5, imgsz=320, verbose=False)
+            print("⚠️ Using regular YOLO - detecting all objects, not just the specified prompt")
+    except Exception as e:
+        print(f"Error in YOLO prediction: {e}")
+        return np.array([])
+    
     boxes = []
     for result in results:
         if result.boxes is not None and len(result.boxes.xyxy) > 0:
@@ -20,8 +56,10 @@ def detect_motorcycles(first_frame, prompt="motorcycle"):
     
     if len(boxes) > 0:
         boxes = np.vstack(boxes)
+        print(f"Detected {len(boxes)} objects")
     else:
         boxes = np.array([])
+        print("No objects detected")
     return boxes
 
 def segment_and_highlight_video(video_path, prompt="motorcycle", highlight_color="red"):
@@ -131,8 +169,8 @@ iface = gr.Interface(
         gr.Dropdown(choices=["red", "green", "blue"], label="Highlight Color", value="red")
     ],
     outputs=gr.Video(label="Highlighted Video"),
-    title="Video Segmentation with MobileSAM and YOLO-World (CPU Optimized)",
-    description="Upload a short video (5-10 seconds), specify a text prompt (e.g., 'motorcycle'), and choose a highlight color. Optimized for CPU with 320x180 resolution.",
+    title="Video Segmentation with MobileSAM and YOLO (CPU Optimized)",
+    description="Upload a short video (5-10 seconds), specify a text prompt (e.g., 'motorcycle'), and choose a highlight color. Uses MobileSAM + YOLO for CPU processing at 320x180 resolution.",
     examples=[
         [None, "motorcycle", "red"],
         [None, "car", "green"],