Update utils/video_processing.py

utils/video_processing.py

@@ -3,6 +3,7 @@ import cv2
 import numpy as np
 from ultralytics import YOLO
 import os
+import torch
 
 # Path to the YOLO model
 MODEL_PATH = 'models/yolov8_model.pt'
@@ -11,11 +12,12 @@ MODEL_PATH = 'models/yolov8_model.pt'
 if not os.path.exists(MODEL_PATH):
     raise FileNotFoundError(f"YOLO model file not found at {MODEL_PATH}. Please ensure 'yolov8_model.pt' is in the 'models/' directory.")
 
-# Load YOLO model
+# Load YOLO model with weights_only=False for compatibility with PyTorch 2.6
 try:
-    model = YOLO(MODEL_PATH)
+    # Explicitly set weights_only=False to allow loading Ultralytics model metadata
+    model = YOLO(MODEL_PATH, weights_only=False)
 except Exception as e:
-    raise RuntimeError(f"Failed to load YOLO model from {MODEL_PATH}: {str(e)}")
+    raise RuntimeError(f"Failed to load YOLO model from {MODEL_PATH}: {str(e)}. Ensure the model is a valid YOLOv8 .pt file from a trusted source.")
 
 def track_ball(video_path: str) -> list:
     """
@@ -87,9 +89,4 @@ def generate_replay(video_path: str, trajectory: list, decision: str) -> str:
                 cv2.line(frame, (int(trajectory[i-1][0]), int(trajectory[i-1][1])),
                          (int(trajectory[i][0]), int(trajectory[i][1])), (255, 0, 0), 2)
         cv2.putText(frame, f"Decision: {decision}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
-        out.write(frame)
-        frame_idx += 1
-
-    cap.release()
-    out.release()
-    return replay_path
+        out