Downgrade PyTorch to 2.0.x for YOLOv5 compatibility and simplify model loading accordingly

app/models/yolo.py

@@ -54,37 +54,9 @@ class MarineSpeciesYOLO:
 
             logger.info(f"Loading YOLOv5 model from: {self.model_path}")
 
-            # Handle PyTorch 2.6+ weights_only issue
-            # Add safe globals for numpy operations that YOLOv5 needs
-            safe_globals = [
-                'numpy.core.multiarray._reconstruct',
-                'numpy.ndarray',
-                'numpy.dtype',
-                'numpy.core.multiarray.scalar',
-                'collections.OrderedDict',
-                'torch._utils._rebuild_tensor_v2',
-                'torch.nn.modules.conv.Conv2d',
-                'torch.nn.modules.batchnorm.BatchNorm2d',
-                'torch.nn.modules.activation.SiLU'
-            ]
-
-            try:
-                # Add safe globals
-                torch.serialization.add_safe_globals(safe_globals)
-                logger.info("Added safe globals for PyTorch loading")
-            except Exception as e:
-                logger.warning(f"Could not add safe globals: {e}")
-
-            # Try loading with safe globals context manager as well
-            try:
-                with torch.serialization.safe_globals(safe_globals):
-                    self.model = yolov5.load(self.model_path, device=self.device)
-                logger.info("Model loaded successfully with safe globals context")
-            except Exception as e:
-                logger.warning(f"Safe globals context failed: {e}")
-                # Fallback to regular loading
-                logger.info("Attempting fallback model loading...")
-                self.model = yolov5.load(self.model_path, device=self.device)
+            # Load model exactly like the original Gradio implementation
+            # Using PyTorch 2.0.1 which doesn't have the weights_only=True issue
+            self.model = yolov5.load(self.model_path, device=self.device)
 
             # Get class names if available
             if hasattr(self.model, 'names'):

requirements.txt

@@ -3,8 +3,8 @@ fastapi==0.104.1
 uvicorn[standard]==0.24.0
 
 # Machine Learning and Computer Vision
-torch>=1.13.0
-torchvision>=0.14.0
+torch==2.0.1
+torchvision==0.15.2
 yolov5==6.2.3
 opencv-python-headless==4.8.1.78
 Pillow==10.1.0