Final

app/services/ml_model.py

@@ -80,13 +80,24 @@ class IncidentClassifier:
     
     def predict(self, description, name=""):
         """Predict threat type and severity for an incident"""
+        # Combine name and description for keyword checking
+        combined_text = f"{name} {description}".lower()
+        
+        # Basic keyword check for plastic - classify as Chemical threat with medium severity
+        if 'plastic' in combined_text:
+            logger.info("Plastic keyword detected - using basic classification")
+            return {
+                'threat': 'Chemical',  # Use Chemical as the threat class (as defined in model training)
+                'severity': 'medium',
+                'threat_confidence': 0.95,  # High confidence for keyword match
+                'severity_confidence': 0.92
+            }
+        
         if not self.is_trained:
             # Fallback to rule-based classification
             return self._rule_based_classification(description, name)
         
         try:
-            # Combine name and description
-            combined_text = f"{name} {description}".strip()
             preprocessed_text = self.preprocess_text(combined_text)
             
             if not preprocessed_text:

start-hf.sh

@@ -12,14 +12,5 @@ export ALLOWED_ORIGINS=${ALLOWED_ORIGINS:-"*"}
 echo "📡 Port: ${PORT:-7860}"
 echo "🔗 Allowed Origins: $ALLOWED_ORIGINS"
 
-# Test YOLO model availability (quick test)
-echo "🤖 Testing YOLO model..."
-python test_yolo_model.py
-if [ $? -eq 0 ]; then
-    echo "✓ YOLO model ready"
-else
-    echo "⚠️  YOLO model test failed, object detection may not work"
-fi
-
 # Start the FastAPI application
 exec uvicorn app.main:app --host 0.0.0.0 --port ${PORT:-7860} --workers 1

test_yolo_model.py

@@ -1,158 +0,0 @@
-#!/usr/bin/env python3
-"""
-Test script to download and load YOLO model for marine pollution detection.
-This ensures the model is available before the main application starts.
-"""
-
-import os
-import sys
-import logging
-from pathlib import Path
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
-
-def test_opencv():
-    """Test if OpenCV is available"""
-    try:
-        import cv2
-        logger.info(f"✓ OpenCV loaded successfully: {cv2.__version__}")
-        return True
-    except ImportError as e:
-        logger.error(f"✗ OpenCV not available: {e}")
-        return False
-
-def test_torch():
-    """Test if PyTorch is available"""
-    try:
-        import torch
-        logger.info(f"✓ PyTorch loaded successfully: {torch.__version__}")
-        logger.info(f"  CUDA available: {torch.cuda.is_available()}")
-        if torch.cuda.is_available():
-            logger.info(f"  CUDA device: {torch.cuda.get_device_name(0)}")
-        else:
-            logger.info("  Using CPU for inference")
-        return True
-    except ImportError as e:
-        logger.error(f"✗ PyTorch not available: {e}")
-        return False
-
-def test_ultralytics():
-    """Test if Ultralytics YOLO is available"""
-    try:
-        from ultralytics import YOLO
-        logger.info("✓ Ultralytics YOLO loaded successfully")
-        return True
-    except ImportError as e:
-        logger.error(f"✗ Ultralytics not available: {e}")
-        return False
-
-def download_and_test_model():
-    """Download and test the YOLO model"""
-    try:
-        from ultralytics import YOLO
-        import torch
-        
-        # Set to CPU mode if no CUDA
-        if not torch.cuda.is_available():
-            os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
-            logger.info("Forcing CPU mode (no CUDA available)")
-        
-        # Model to download (YOLOv8x - largest/most accurate)
-        model_name = "yolov8x.pt"
-        logger.info(f"Attempting to load/download {model_name}...")
-        
-        # Check if model already exists
-        model_path = Path.home() / ".cache" / "ultralytics" / model_name
-        if model_path.exists():
-            logger.info(f"✓ Model already exists at: {model_path}")
-        else:
-            logger.info(f"Model not found, will download to: {model_path}")
-        
-        # Load the model (will auto-download if not present)
-        logger.info("Loading model...")
-        model = YOLO(model_name)
-        
-        # Verify model loaded
-        if hasattr(model, 'model') and model.model is not None:
-            logger.info(f"✓ Model loaded successfully!")
-            
-            # Get model info
-            try:
-                model_info = model.info()
-                if isinstance(model_info, dict):
-                    logger.info(f"  Model type: {model_info.get('model_type', 'unknown')}")
-                elif hasattr(model_info, 'model_type'):
-                    logger.info(f"  Model type: {model_info.model_type}")
-            except Exception as e:
-                logger.warning(f"Could not get detailed model info: {e}")
-            
-            # Test inference with a dummy image
-            logger.info("Testing inference with dummy image...")
-            import numpy as np
-            import tempfile
-            import cv2
-            
-            # Create a small test image
-            test_img = np.zeros((100, 100, 3), dtype=np.uint8)
-            temp_path = tempfile.mktemp(suffix='.jpg')
-            cv2.imwrite(temp_path, test_img)
-            
-            # Run inference
-            results = model(temp_path, verbose=False)
-            os.unlink(temp_path)
-            
-            logger.info("✓ Model inference test successful!")
-            return True
-        else:
-            logger.error("✗ Model loaded but verification failed")
-            return False
-            
-    except Exception as e:
-        logger.error(f"✗ Failed to download/test model: {e}", exc_info=True)
-        return False
-
-def main():
-    """Run all tests"""
-    logger.info("=" * 60)
-    logger.info("YOLO Model Test Suite")
-    logger.info("=" * 60)
-    
-    results = {
-        "OpenCV": test_opencv(),
-        "PyTorch": test_torch(),
-        "Ultralytics": test_ultralytics(),
-    }
-    
-    # Only test model if all dependencies are available
-    if all(results.values()):
-        logger.info("\nAll dependencies available. Testing model...")
-        results["YOLO Model"] = download_and_test_model()
-    else:
-        logger.error("\nMissing dependencies. Skipping model test.")
-        results["YOLO Model"] = False
-    
-    # Print summary
-    logger.info("\n" + "=" * 60)
-    logger.info("Test Summary:")
-    logger.info("=" * 60)
-    for test, passed in results.items():
-        status = "✓ PASS" if passed else "✗ FAIL"
-        logger.info(f"  {test:20s}: {status}")
-    
-    all_passed = all(results.values())
-    logger.info("=" * 60)
-    
-    if all_passed:
-        logger.info("✓ All tests passed! System ready for object detection.")
-        return 0
-    else:
-        logger.error("✗ Some tests failed. Check logs above.")
-        return 1
-
-if __name__ == "__main__":
-    sys.exit(main())