Update utils/video_processing.py

utils/video_processing.py

@@ -12,12 +12,24 @@ 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 with weights_only=False for compatibility with PyTorch 2.6
+# Load YOLO model
 try:
-    # Explicitly set weights_only=False to allow loading Ultralytics model metadata
-    model = YOLO(MODEL_PATH, weights_only=False)
+    # Load the model using Ultralytics YOLO
+    model = YOLO(MODEL_PATH)
 except Exception as 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.")
+    # If loading fails due to weights_only issue, try manual loading
+    try:
+        # Manually load the checkpoint with weights_only=False
+        checkpoint = torch.load(MODEL_PATH, map_location='cpu', weights_only=False)
+        model = YOLO('yolov8n.yaml')  # Load model architecture from YAML
+        model.load_state_dict(checkpoint['model'].state_dict())  # Load weights
+    except Exception as inner_e:
+        raise RuntimeError(
+            f"Failed to load YOLO model from {MODEL_PATH}: {str(e)}. "
+            f"Manual loading also failed: {str(inner_e)}. "
+            "Ensure the model is a valid YOLOv8 .pt file from a trusted source. "
+            "You may need to re-save the model or use a pre-trained model like yolov8n.pt."
+        )
 
 def track_ball(video_path: str) -> list:
     """
@@ -89,4 +101,9 @@ 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
+        out.write(frame)
+        frame_idx += 1
+
+    cap.release()
+    out.release()
+    return replay_path