Downgrade YOLOv5 for compatibility, simplify model loading, fix Pydantic warnings, and improve error handling

Dockerfile

@@ -17,51 +17,31 @@ RUN pip install huggingface_hub
 # Create models directory
 RUN mkdir -p /models
 
-# Create a Python script for model downloading
-RUN echo 'from huggingface_hub import hf_hub_download\n\
-import os\n\
-\n\
-def download_model():\n\
-    try:\n\
-        print("Attempting to download model...")\n\
-        hf_hub_download(\n\
-            "seamo-ai/marina-species-v1",\n\
-            "marina-benthic-33k.pt",\n\
-            local_dir="/models",\n\
-            local_dir_use_symlinks=False\n\
-        )\n\
-        print("Model downloaded successfully")\n\
-        return True\n\
-    except Exception as e:\n\
-        print(f"Model download failed: {e}")\n\
-        # Create a placeholder file to prevent build failure\n\
-        with open("/models/marina-benthic-33k.pt", "w") as f:\n\
-            f.write("placeholder")\n\
-        print("Created placeholder model file")\n\
-        return False\n\
-\n\
-def download_names():\n\
-    try:\n\
-        print("Attempting to download class names...")\n\
-        hf_hub_download(\n\
-            "seamo-ai/marina-species-v1",\n\
-            "marina-benthic-33k.names",\n\
-            local_dir="/models",\n\
-            local_dir_use_symlinks=False\n\
-        )\n\
-        print("Class names downloaded successfully")\n\
-        return True\n\
-    except Exception as e:\n\
-        print(f"Class names download failed: {e}")\n\
-        return False\n\
-\n\
-if __name__ == "__main__":\n\
-    download_model()\n\
-    download_names()\n\
-' > /tmp/download_models.py
-
-# Run the model download script
-RUN python /tmp/download_models.py
+# Download models from HuggingFace Hub
+RUN python -c "\
+from huggingface_hub import hf_hub_download; \
+import os; \
+try: \
+    print('Downloading model...'); \
+    hf_hub_download('seamo-ai/marina-species-v1', 'marina-benthic-33k.pt', local_dir='/models', local_dir_use_symlinks=False); \
+    print('Model downloaded successfully'); \
+except Exception as e: \
+    print(f'Model download failed: {e}'); \
+    with open('/models/marina-benthic-33k.pt', 'w') as f: \
+        f.write('placeholder'); \
+    print('Created placeholder model file'); \
+"
+
+# Download class names file
+RUN python -c "\
+from huggingface_hub import hf_hub_download; \
+try: \
+    print('Downloading class names...'); \
+    hf_hub_download('seamo-ai/marina-species-v1', 'marina-benthic-33k.names', local_dir='/models', local_dir_use_symlinks=False); \
+    print('Class names downloaded successfully'); \
+except Exception as e: \
+    print(f'Class names download failed: {e}'); \
+"
 
 # Stage 2: Build the application
 FROM python:3.10-slim

app/models/inference.py

@@ -77,15 +77,17 @@ class InferenceRequest(BaseModel):
 
 class InferenceResponse(BaseModel):
     """Response model for marine species detection."""
+    model_config = {"protected_namespaces": ()}
+
     detections: List[Detection] = Field(..., description="List of detected marine species")
     annotated_image: Optional[str] = Field(
-        default=None, 
+        default=None,
         description="Base64 encoded annotated image (if requested)"
     )
     processing_time: float = Field(..., description="Processing time in seconds")
     model_info: ModelInfo = Field(..., description="Information about the model used")
     image_dimensions: Dict[str, int] = Field(
-        ..., 
+        ...,
         description="Original image dimensions (width, height)"
     )
 

app/models/yolo.py

@@ -54,24 +54,8 @@ class MarineSpeciesYOLO:
 
             logger.info(f"Loading YOLOv5 model from: {self.model_path}")
 
-            # Handle PyTorch 2.6+ weights_only issue
-            import torch
-            import warnings
-
-            # Temporarily suppress warnings and set safe globals for numpy
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-
-                # Add safe globals for numpy operations that YOLOv5 needs
-                torch.serialization.add_safe_globals([
-                    'numpy.core.multiarray._reconstruct',
-                    'numpy.ndarray',
-                    'numpy.dtype',
-                    'numpy.core.multiarray.scalar',
-                ])
-
-                # Load the model with YOLOv5
-                self.model = yolov5.load(self.model_path, device=self.device)
+            # Load model exactly like the original Gradio implementation
+            self.model = yolov5.load(self.model_path, device=self.device)
 
             # Get class names if available
             if hasattr(self.model, 'names'):

requirements.txt

@@ -5,7 +5,7 @@ uvicorn[standard]==0.24.0
 # Machine Learning and Computer Vision
 torch>=1.13.0
 torchvision>=0.14.0
-yolov5==7.0.13
+yolov5==6.2.3
 opencv-python-headless==4.8.1.78
 Pillow==10.1.0
 numpy==1.24.3