Update models/yolov8_model.pt

models/yolov8_model.pt

@@ -1,6 +1,9 @@
 # scripts/download_yolov8_model.py
 from ultralytics import YOLO
 import os
+import torch
+import requests
+import shutil
 
 # Directory to save the model
 MODEL_DIR = "models"
@@ -9,14 +12,49 @@ os.makedirs(MODEL_DIR, exist_ok=True)
 # Path to save the model
 MODEL_PATH = os.path.join(MODEL_DIR, "yolov8_model.pt")
 
-# Download pre-trained YOLOv8 nano model
+# URL for yolov8n.pt from Ultralytics GitHub releases
+MODEL_URL = "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt"
+
+def download_model(url, output_path):
+    """Download model file from URL and save to output_path."""
+    print(f"Downloading model from {url}...")
+    try:
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(output_path, 'wb') as f:
+                shutil.copyfileobj(response.raw, f)
+            print(f"Model downloaded to {output_path}")
+        else:
+            raise Exception(f"Failed to download model: HTTP {response.status_code}")
+    except Exception as e:
+        raise Exception(f"Download error: {str(e)}")
+
+def verify_model(path):
+    """Verify that the model file is a valid PyTorch checkpoint."""
+    print(f"Verifying model at {path}...")
+    try:
+        checkpoint = torch.load(path, map_location='cpu', weights_only=False)
+        print(f"Model verified as valid PyTorch checkpoint")
+        return True
+    except Exception as e:
+        raise Exception(f"Verification failed: {str(e)}")
+
+# Download and verify model
 try:
-    model = YOLO("yolov8n.pt")  # Automatically downloads from Ultralytics
-    model.save(MODEL_PATH)
-    print(f"Model saved to {MODEL_PATH}")
-    # Verify file integrity
-    import torch
-    checkpoint = torch.load(MODEL_PATH, map_location='cpu', weights_only=False)
-    print(f"Model file verified as valid PyTorch checkpoint")
+    # Remove existing file to avoid using corrupted version
+    if os.path.exists(MODEL_PATH):
+        os.remove(MODEL_PATH)
+        print(f"Removed existing file at {MODEL_PATH}")
+
+    # Download yolov8n.pt
+    download_model(MODEL_URL, MODEL_PATH)
+
+    # Verify the downloaded file
+    verify_model(MODEL_PATH)
+
+    # Load with YOLO to ensure compatibility
+    model = YOLO(MODEL_PATH)
+    print(f"Model successfully loaded with Ultralytics YOLO")
 except Exception as e:
-    print(f"Failed to download, save, or verify model: {str(e)}")
+    print(f"Failed to download or verify model: {str(e)}")
+    print("Please manually download yolov8n.pt from https://github.com/ultralytics/assets/releases and place it in backend/models/yolov8_model.pt")