Update main.py

main.py

@@ -1,7 +1,7 @@
 """
 FastAPI Backend for Respiratory Symptom Analysis
+Updated with proper model weight loading and health classification system
 Deployed on HuggingFace Spaces for use with Netlify frontend
-Fixed for model loading compatibility issues
 """
 
 from fastapi import FastAPI, File, UploadFile, HTTPException
@@ -138,7 +138,7 @@ class PurePyTorchInferenceModel(nn.Module):
 app = FastAPI(
     title="🫁 Respiratory Symptom Analysis API",
     description="AI-powered respiratory symptom detection from cough audio",
-    version="2.0.0",
+    version="2.1.0",
     docs_url="/docs",
     redoc_url="/redoc"
 )
@@ -154,7 +154,7 @@ app.add_middleware(
 
 class RespiratoryAnalysisService:
     """
-    Service class for respiratory symptom analysis
+    Enhanced service class for respiratory symptom analysis with proper model loading
     """
     
     def __init__(self, config_path: str = "optimized_model_cpu/model_config.json"):
@@ -163,6 +163,8 @@ class RespiratoryAnalysisService:
         self.model = None
         self.config = None
         self.preprocessor = None
+        self.weights_loaded = False  # Track if real weights are loaded
+        self.neutral_threshold = 0.35  # Below this = neutral/healthy
         
         # Load configuration and model
         self.load_config()
@@ -196,7 +198,7 @@ class RespiratoryAnalysisService:
                         'fever': '#FF6B6B', 'cold': '#4ECDC4', 'sorethroat': '#45B7D1',
                         'lossofsmell': '#96CEB4', 'fatigue': '#FFEAA7', 'cough': '#DDA0DD'
                     },
-                    'model_version': '2.0',
+                    'model_version': '2.1',
                     'optimization_settings': {'torch_threads': 4}
                 }
                 print("⚠️ Using default configuration")
@@ -205,7 +207,7 @@ class RespiratoryAnalysisService:
             raise RuntimeError(f"Failed to load config: {str(e)}")
     
     def create_and_load_model(self):
-        """Create model and try to load weights from available files"""
+        """Create model and try to load weights from available files with priority order"""
         try:
             # Create model with correct architecture
             self.model = PurePyTorchInferenceModel(
@@ -213,54 +215,57 @@ class RespiratoryAnalysisService:
                 confidence_thresholds=self.config['confidence_thresholds']
             )
             
-            # Try to load state dict from available files
-            state_dict_files = [
-                "optimized_model_cpu/model_pytorch_state_dict.pt",
-                "optimized_model_cpu/model_quantized_state_dict.pt"
+            print("🔍 Searching for model weight files...")
+            
+            # ✅ PRIORITY ORDER: Try different model files with detailed logging
+            weight_files_to_try = [
+                # Highest priority - state dicts (most compatible)
+                ("optimized_model_cpu/model_pytorch_state_dict.pt", "PyTorch State Dict", "state_dict"),
+                ("optimized_model_cpu/model_quantized_state_dict.pt", "Quantized State Dict", "state_dict"),
+                
+                # Medium priority - full models
+                ("optimized_model_cpu/model_pytorch.pt", "Full PyTorch Model", "full_model"),
+                ("optimized_model_cpu/model_quantized.pt", "Quantized PyTorch Model", "full_model"),
+                
+                # Lower priority - TorchScript (compatibility issues)
+                ("optimized_model_cpu/model_torchscript.pt", "TorchScript Model", "torchscript"),
             ]
             
-            model_loaded = False
-            for state_dict_file in state_dict_files:
-                if Path(state_dict_file).exists():
+            for weight_file, model_type, load_type in weight_files_to_try:
+                if Path(weight_file).exists():
+                    file_size = Path(weight_file).stat().st_size / (1024*1024)  # Size in MB
+                    print(f"📁 Found {model_type}: {weight_file} ({file_size:.1f}MB)")
+                    
                     try:
-                        state_dict = torch.load(state_dict_file, map_location='cpu')
-                        
-                        # Handle different state dict formats
-                        if 'state_dict' in state_dict:
-                            state_dict = state_dict['state_dict']
-                        
-                        # Filter out incompatible keys (like attention layers)
-                        filtered_state_dict = {}
-                        for key, value in state_dict.items():
-                            # Skip attention-related keys that might cause issues
-                            if 'symptom_attention' in key:
-                                continue
-                            # Skip COVID classifier if present
-                            if 'covid_classifier' in key:
-                                continue
-                            filtered_state_dict[key] = value
-                        
-                        # Load with strict=False to ignore missing/extra keys
-                        missing_keys, unexpected_keys = self.model.load_state_dict(filtered_state_dict, strict=False)
-                        
-                        if missing_keys:
-                            print(f"⚠️ Missing keys (will use random initialization): {missing_keys}")
-                        if unexpected_keys:
-                            print(f"⚠️ Unexpected keys (ignored): {unexpected_keys}")
-                        
-                        print(f"✅ Model weights loaded from {state_dict_file}")
-                        model_loaded = True
-                        break
-                        
+                        if load_type == "state_dict":
+                            success = self._load_state_dict(weight_file, model_type)
+                        elif load_type == "full_model":
+                            success = self._load_full_model(weight_file, model_type)
+                        elif load_type == "torchscript":
+                            success = self._load_torchscript_model(weight_file, model_type)
+                        else:
+                            success = False
+                            
+                        if success:
+                            self.weights_loaded = True
+                            print(f"✅ Successfully loaded {model_type}")
+                            break
+                            
                     except Exception as e:
-                        print(f"⚠️ Failed to load {state_dict_file}: {str(e)}")
+                        print(f"⚠️ Failed to load {model_type}: {str(e)}")
                         continue
                 else:
-                    print(f"⚠️ State dict file not found: {state_dict_file}")
+                    print(f"❌ Not found: {weight_file}")
             
-            if not model_loaded:
-                print("⚠️ No compatible model weights found, using random initialization")
-                print("⚠️ Model will give random predictions - this is for testing only!")
+            if not self.weights_loaded:
+                print("\n❌ WARNING: Using random model weights!")
+                print("❌ All predictions will be random (~50% confidence)")
+                print("❌ Please check your model files in optimized_model_cpu/")
+                print("❌ Expected files:")
+                for file_path, _, _ in weight_files_to_try:
+                    print(f"   - {file_path}")
+            else:
+                print(f"✅ Model ready with trained weights")
             
             # Set model to evaluation mode
             self.model.eval()
@@ -271,13 +276,107 @@ class RespiratoryAnalysisService:
         except Exception as e:
             raise RuntimeError(f"Failed to create/load model: {str(e)}")
     
+    def _load_state_dict(self, weight_file: str, model_type: str) -> bool:
+        """Load model from state dict file"""
+        try:
+            checkpoint = torch.load(weight_file, map_location='cpu')
+            
+            # Handle different checkpoint formats
+            if isinstance(checkpoint, dict):
+                if 'state_dict' in checkpoint:
+                    state_dict = checkpoint['state_dict']
+                elif 'model_state_dict' in checkpoint:
+                    state_dict = checkpoint['model_state_dict']
+                else:
+                    state_dict = checkpoint
+            else:
+                state_dict = checkpoint
+            
+            # Remove any incompatible keys
+            filtered_state_dict = {}
+            for key, value in state_dict.items():
+                # Skip keys that might cause issues
+                if any(skip in key for skip in ['symptom_attention', 'covid_classifier', 'aux_']):
+                    print(f"   Skipping incompatible key: {key}")
+                    continue
+                filtered_state_dict[key] = value
+            
+            # Load weights
+            missing_keys, unexpected_keys = self.model.load_state_dict(filtered_state_dict, strict=False)
+            
+            # Check if enough weights were loaded
+            loaded_keys = len(filtered_state_dict) - len(missing_keys)
+            total_keys = len(self.model.state_dict())
+            load_percentage = (loaded_keys / total_keys) * 100
+            
+            print(f"   📊 Loaded {loaded_keys}/{total_keys} parameters ({load_percentage:.1f}%)")
+            
+            if missing_keys:
+                print(f"   ⚠️ Missing keys: {len(missing_keys)} (using random initialization)")
+            if unexpected_keys:
+                print(f"   ⚠️ Unexpected keys: {len(unexpected_keys)} (ignored)")
+            
+            # Consider successful if we loaded most parameters
+            return load_percentage > 50
+            
+        except Exception as e:
+            print(f"   ❌ State dict loading failed: {str(e)}")
+            return False
+    
+    def _load_full_model(self, weight_file: str, model_type: str) -> bool:
+        """Load full model file"""
+        try:
+            loaded_model = torch.load(weight_file, map_location='cpu')
+            
+            if hasattr(loaded_model, 'state_dict'):
+                # Extract state dict from full model
+                state_dict = loaded_model.state_dict()
+                return self._load_state_dict_direct(state_dict)
+            else:
+                # Try to use as state dict directly
+                return self._load_state_dict_direct(loaded_model)
+                
+        except Exception as e:
+            print(f"   ❌ Full model loading failed: {str(e)}")
+            return False
+    
+    def _load_torchscript_model(self, weight_file: str, model_type: str) -> bool:
+        """Load TorchScript model (with known compatibility issues)"""
+        try:
+            scripted_model = torch.jit.load(weight_file, map_location='cpu')
+            scripted_model.eval()
+            
+            # Replace the model entirely with TorchScript version
+            self.model = scripted_model
+            print(f"   ✅ Using TorchScript model directly")
+            return True
+            
+        except Exception as e:
+            print(f"   ❌ TorchScript loading failed: {str(e)}")
+            return False
+    
+    def _load_state_dict_direct(self, state_dict: Dict) -> bool:
+        """Helper to load state dict directly"""
+        try:
+            missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
+            loaded_keys = len(state_dict) - len(missing_keys)
+            total_keys = len(self.model.state_dict())
+            load_percentage = (loaded_keys / total_keys) * 100
+            
+            print(f"   📊 Loaded {loaded_keys}/{total_keys} parameters ({load_percentage:.1f}%)")
+            return load_percentage > 50
+            
+        except Exception as e:
+            print(f"   ❌ Direct state dict loading failed: {str(e)}")
+            return False
+    
     def setup_preprocessor(self):
         """Initialize audio preprocessor"""
         self.preprocessor = RespiratoryAudioPreprocessor()
         print("✅ Audio preprocessor initialized")
     
     def predict_symptoms(self, audio_file_path: str) -> Dict[str, Any]:
-        """Predict respiratory symptoms from audio file"""
+        """Predict respiratory symptoms with enhanced threshold logic and health classification"""
         try:
             start_time = time.time()
             
@@ -289,22 +388,60 @@ class RespiratoryAnalysisService:
             inference_start = time.time()
             with torch.no_grad():
                 outputs = self.model(tensor_input)
-                
             inference_time = time.time() - inference_start
             
             # Parse outputs
             probabilities = outputs['probabilities'].squeeze().numpy()
-            predictions = outputs['predictions'].squeeze().numpy()
             
-            # Format results
-            results = self.format_results(probabilities, predictions)
+            # ✅ ENHANCED THRESHOLD LOGIC with neutral detection
+            detected_symptoms = []
+            
+            for i, symptom in enumerate(self.config['target_symptoms']):
+                prob = float(probabilities[i])
+                symptom_threshold = self.config['confidence_thresholds'][symptom]
+                
+                # Apply dual threshold system:
+                # 1. Must be above symptom-specific threshold
+                # 2. Must be above neutral threshold to avoid false positives
+                effective_threshold = max(symptom_threshold, self.neutral_threshold)
+                is_detected = prob >= effective_threshold
+                
+                if is_detected:
+                    detected_symptoms.append({
+                        'symptom': symptom,
+                        'display_name': self.config['symptom_display_names'][symptom],
+                        'confidence': prob,
+                        'color': self.config['symptom_colors'][symptom],
+                        'threshold_used': effective_threshold
+                    })
+            
+            # ✅ DETERMINE OVERALL HEALTH STATUS
+            max_confidence = np.max(probabilities)
+            
+            if not detected_symptoms:
+                if max_confidence < self.neutral_threshold:
+                    health_status = "healthy"
+                    status_message = "No symptoms detected - appears healthy"
+                else:
+                    health_status = "inconclusive" 
+                    status_message = "Some patterns detected but below confidence threshold"
+            else:
+                health_status = "symptoms_detected"
+                status_message = f"{len(detected_symptoms)} symptom(s) detected"
+            
+            # Format results with enhanced health classification
+            results = self.format_results_enhanced(
+                probabilities, detected_symptoms, health_status, status_message, max_confidence
+            )
             
-            # Add timing info
+            # Add comprehensive processing info
             results['processing_info'] = {
                 'preprocessing_time_ms': round(preprocessing_time * 1000, 1),
                 'inference_time_ms': round(inference_time * 1000, 1),
                 'total_time_ms': round((preprocessing_time + inference_time) * 1000, 1),
-                'model_status': 'loaded' if hasattr(self, '_weights_loaded') else 'random_weights'
+                'model_weights_loaded': self.weights_loaded,
+                'neutral_threshold': self.neutral_threshold,
+                'max_confidence': round(max_confidence, 3)
             }
             
             return results
@@ -312,176 +449,233 @@ class RespiratoryAnalysisService:
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")
     
-    def format_results(self, probabilities: np.ndarray, predictions: np.ndarray) -> Dict[str, Any]:
-        """Format prediction results for API response"""
+    def format_results_enhanced(self, probabilities, detected_symptoms, health_status, status_message, max_confidence):
+        """Enhanced results formatting with health classification"""
+        
         results = {
-            'detected_symptoms': [],
+            'detected_symptoms': detected_symptoms,
             'all_symptoms': {},
             'summary': {},
-            'recommendations': []
+            'recommendations': [],
+            'health_classification': health_status
         }
         
-        # Process each symptom
+        # Process all symptoms with enhanced threshold information
         for i, symptom in enumerate(self.config['target_symptoms']):
             prob = float(probabilities[i])
-            pred = bool(predictions[i])
-            display_name = self.config['symptom_display_names'][symptom]
-            threshold = self.config['confidence_thresholds'][symptom]
+            original_threshold = self.config['confidence_thresholds'][symptom]
+            effective_threshold = max(original_threshold, self.neutral_threshold)
+            detected = prob >= effective_threshold
             
-            # All symptoms with details
             results['all_symptoms'][symptom] = {
-                'display_name': display_name,
+                'display_name': self.config['symptom_display_names'][symptom],
                 'confidence': prob,
-                'detected': pred,
-                'threshold': threshold,
+                'detected': detected,
+                'original_threshold': original_threshold,
+                'effective_threshold': effective_threshold,
+                'neutral_threshold': self.neutral_threshold,
                 'color': self.config['symptom_colors'][symptom]
             }
-            
-            # Detected symptoms only
-            if pred:
-                results['detected_symptoms'].append({
-                    'symptom': symptom,
-                    'display_name': display_name,
-                    'confidence': prob,
-                    'color': self.config['symptom_colors'][symptom]
-                })
-        
-        # Sort detected symptoms by confidence
-        results['detected_symptoms'].sort(key=lambda x: x['confidence'], reverse=True)
         
-        # Generate summary
+        # Enhanced summary with health classification
         results['summary'] = {
-            'total_detected': len(results['detected_symptoms']),
-            'highest_confidence': results['detected_symptoms'][0]['confidence'] if results['detected_symptoms'] else 0.0,
-            'status': 'symptoms_detected' if results['detected_symptoms'] else 'no_symptoms'
+            'total_detected': len(detected_symptoms),
+            'highest_confidence': max([s['confidence'] for s in detected_symptoms], default=0.0),
+            'max_overall_confidence': max_confidence,
+            'status': health_status,
+            'status_message': status_message,
+            'neutral_threshold': self.neutral_threshold,
+            'weights_status': 'trained' if self.weights_loaded else 'random'
         }
         
-        # Generate recommendations
-        if len(results['detected_symptoms']) == 0:
+        # ✅ ENHANCED RECOMMENDATIONS based on health status
+        if health_status == "healthy":
+            results['recommendations'] = [
+                "✅ No significant respiratory symptoms detected",
+                "Your cough patterns appear normal and healthy",
+                "Continue maintaining good respiratory health practices",
+                "This screening is for informational purposes only"
+            ]
+        elif health_status == "inconclusive":
             results['recommendations'] = [
-                "No significant respiratory symptoms detected.",
-                "Continue monitoring your health.",
-                "This screening is for informational purposes only."
+                "⚠️ Some respiratory patterns detected but below confidence threshold",
+                "Consider monitoring your symptoms over the next few days", 
+                "If symptoms persist or worsen, consult a healthcare provider",
+                "This AI screening should not replace professional medical advice"
             ]
-        elif len(results['detected_symptoms']) == 1:
-            symptom_name = results['detected_symptoms'][0]['display_name']
+        elif len(detected_symptoms) == 1:
+            symptom_name = detected_symptoms[0]['display_name']
+            confidence = detected_symptoms[0]['confidence']
             results['recommendations'] = [
-                f"Detected: {symptom_name}",
-                "Consider monitoring symptoms and consult healthcare provider if symptoms persist.",
-                "This AI screening should not replace professional medical advice."
+                f"🔍 Detected: {symptom_name} (confidence: {confidence:.1%})",
+                "Monitor this symptom and note any changes or progression",
+                "Consider consulting a healthcare provider if symptoms persist or worsen",
+                "This AI screening should not replace professional medical advice"
             ]
         else:
-            symptom_names = [s['display_name'] for s in results['detected_symptoms']]
+            symptom_names = [s['display_name'] for s in detected_symptoms]
             results['recommendations'] = [
-                f"Multiple symptoms detected: {', '.join(symptom_names)}",
-                "Please consult a healthcare provider for proper evaluation.",
-                "This AI screening should not replace professional medical advice."
+                f"🚨 Multiple symptoms detected: {', '.join(symptom_names)}",
+                "Multiple symptoms may indicate a need for medical attention",
+                "Please consult a healthcare provider for proper evaluation and diagnosis",
+                "This AI screening should not replace professional medical advice"
             ]
         
+        # Add model status warning if using random weights
+        if not self.weights_loaded:
+            results['recommendations'].insert(0, 
+                "⚠️ DEVELOPMENT MODE: Model using random weights - results are not medically valid"
+            )
+        
         return results
 
-# Initialize service
-print("🚀 Initializing Respiratory Analysis Service...")
+# Initialize service with enhanced error handling
+print("🚀 Initializing Enhanced Respiratory Analysis Service...")
 try:
     service = RespiratoryAnalysisService()
     print("✅ Service initialized successfully!")
+    print(f"   Model weights loaded: {'Yes' if service.weights_loaded else 'No (using random weights)'}")
+    print(f"   Neutral threshold: {service.neutral_threshold}")
 except Exception as e:
     print(f"❌ Service initialization failed: {str(e)}")
     service = None
 
-# API Routes
+# =================== API ROUTES ===================
+
 @app.get("/")
 async def root():
-    """Root endpoint with API information"""
+    """Root endpoint with enhanced API information"""
     if service is None:
         return {
             "service": "Respiratory Symptom Analysis API",
-            "version": "2.0.0",
+            "version": "2.1.0",
             "status": "error - service not initialized"
         }
     
     return {
-        "service": "Respiratory Symptom Analysis API",
-        "version": "2.0.0",
+        "service": "Respiratory Symptom Analysis API", 
+        "version": "2.1.0",
         "status": "active",
+        "model_status": "trained_weights" if service.weights_loaded else "random_weights",
+        "health_classification": ["healthy", "symptoms_detected", "inconclusive"],
+        "neutral_threshold": service.neutral_threshold,
         "endpoints": {
             "analyze": "/analyze",
-            "health": "/health",
+            "health": "/health", 
             "info": "/info",
             "docs": "/docs"
         },
-        "supported_symptoms": list(service.config['target_symptoms']),
+        "supported_symptoms": service.config['target_symptoms'],
         "model_info": {
             "version": service.config['model_version'],
-            "optimization": "CPU-optimized"
+            "optimization": "CPU-optimized with health classification"
         }
     }
 
 @app.get("/health")
 async def health_check():
-    """Health check endpoint"""
+    """Enhanced health check with detailed model status"""
+    model_files_status = {
+        "pytorch_state_dict": Path("optimized_model_cpu/model_pytorch_state_dict.pt").exists(),
+        "quantized_state_dict": Path("optimized_model_cpu/model_quantized_state_dict.pt").exists(),
+        "pytorch_full": Path("optimized_model_cpu/model_pytorch.pt").exists(),
+        "quantized_full": Path("optimized_model_cpu/model_quantized.pt").exists(),
+        "torchscript": Path("optimized_model_cpu/model_torchscript.pt").exists(),
+        "config": Path("optimized_model_cpu/model_config.json").exists()
+    }
+    
     return {
         "status": "healthy" if service is not None else "unhealthy",
         "timestamp": time.time(),
+        "service_ready": service is not None,
         "model_loaded": service.model is not None if service else False,
         "config_loaded": service.config is not None if service else False,
-        "files_available": {
-            "config": Path("optimized_model_cpu/model_config.json").exists(),
-            "pytorch_state": Path("optimized_model_cpu/model_pytorch_state_dict.pt").exists(),
-            "quantized_state": Path("optimized_model_cpu/model_quantized_state_dict.pt").exists()
-        }
+        "model_weights_status": "trained" if (service and service.weights_loaded) else "random",
+        "neutral_threshold": service.neutral_threshold if service else None,
+        "health_classification_enabled": True,
+        "model_files_available": model_files_status,
+        "files_found": sum(model_files_status.values()),
+        "critical_files_missing": not (model_files_status["config"] and 
+                                     any([model_files_status["pytorch_state_dict"],
+                                          model_files_status["quantized_state_dict"], 
+                                          model_files_status["pytorch_full"]]))
     }
 
 @app.get("/info")
 async def get_info():
-    """Get model and service information"""
+    """Get comprehensive model and service information"""
     if service is None:
         return {"error": "Service not initialized"}
     
     return {
         "model_info": {
-            "version": service.config.get('model_version', '2.0'),
+            "version": service.config.get('model_version', '2.1'),
             "target_symptoms": service.config['target_symptoms'],
             "symptom_display_names": service.config['symptom_display_names'],
-            "confidence_thresholds": service.config['confidence_thresholds']
+            "confidence_thresholds": service.config['confidence_thresholds'],
+            "weights_loaded": service.weights_loaded,
+            "neutral_threshold": service.neutral_threshold,
+            "health_classifications": ["healthy", "symptoms_detected", "inconclusive"]
         },
         "preprocessing_info": service.preprocessor.get_preprocessing_info(),
-        "supported_formats": ["wav", "mp3", "flac", "ogg", "m4a"],
-        "api_version": "2.0.0"
+        "supported_formats": ["wav", "mp3", "flac", "ogg", "m4a", "webm"],
+        "max_duration": "30 seconds",
+        "max_file_size": "10MB",
+        "api_version": "2.1.0",
+        "features": {
+            "health_classification": True,
+            "neutral_detection": True,
+            "dual_threshold_system": True,
+            "trained_weights": service.weights_loaded
+        }
     }
 
 @app.post("/analyze")
 async def analyze_audio(audio_file: UploadFile = File(...)):
-    """Analyze audio file for respiratory symptoms"""
+    """
+    Enhanced audio analysis with health classification
+    
+    Returns:
+    - Detected symptoms with confidence scores
+    - Health classification (healthy/symptoms_detected/inconclusive)  
+    - Enhanced recommendations based on health status
+    - Model weight status for debugging
+    """
     if service is None:
         raise HTTPException(status_code=503, detail="Service not available")
     
-    # Validate file
-    allowed_types = ['audio/wav', 'audio/mpeg', 'audio/flac', 'audio/ogg', 'audio/x-m4a', 'audio/mp4']
+    # Validate file type (including WebM for browser recordings)
+    allowed_types = [
+        'audio/wav', 'audio/mpeg', 'audio/mp3', 'audio/flac', 
+        'audio/ogg', 'audio/x-m4a', 'audio/mp4', 'audio/webm'
+    ]
+    
     if audio_file.content_type not in allowed_types:
         raise HTTPException(
-            status_code=400, 
-            detail=f"Unsupported format: {audio_file.content_type}"
+            status_code=400,
+            detail=f"Unsupported format: {audio_file.content_type}. Supported: {', '.join(allowed_types)}"
         )
     
-    # Validate size
+    # Validate file size
     content = await audio_file.read()
-    if len(content) > 10 * 1024 * 1024:  # 10MB
-        raise HTTPException(status_code=400, detail="File too large")
+    max_size = 10 * 1024 * 1024  # 10MB
+    if len(content) > max_size:
+        raise HTTPException(status_code=400, detail="File too large. Maximum size: 10MB")
     
     try:
-        # Save temporarily
+        # Save uploaded file temporarily  
         file_extension = audio_file.filename.split('.')[-1] if audio_file.filename else 'wav'
         with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}") as temp_file:
             temp_file.write(content)
             temp_file_path = temp_file.name
         
-        # Analyze
+        # Analyze audio with enhanced health classification
         results = service.predict_symptoms(temp_file_path)
         
-        # Cleanup
+        # Clean up temporary file
         os.unlink(temp_file_path)
         
+        # Return enhanced results
         return JSONResponse(
             status_code=200,
             content={
@@ -490,7 +684,9 @@ async def analyze_audio(audio_file: UploadFile = File(...)):
                 "metadata": {
                     "filename": audio_file.filename,
                     "file_size_bytes": len(content),
-                    "timestamp": time.time()
+                    "content_type": audio_file.content_type,
+                    "timestamp": time.time(),
+                    "api_version": "2.1.0"
                 }
             }
         )
@@ -498,13 +694,37 @@ async def analyze_audio(audio_file: UploadFile = File(...)):
     except HTTPException:
         raise
     except Exception as e:
+        # Clean up temporary file if exists
         if 'temp_file_path' in locals():
             try:
                 os.unlink(temp_file_path)
             except:
                 pass
+        
         raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
 
+# Global exception handler
+@app.exception_handler(Exception)
+async def global_exception_handler(request, exc):
+    """Global exception handler with detailed error information"""
+    return JSONResponse(
+        status_code=500,
+        content={
+            "success": False,
+            "error": "Internal server error",
+            "detail": str(exc),
+            "model_status": "trained_weights" if (service and service.weights_loaded) else "random_weights",
+            "timestamp": time.time()
+        }
+    )
+
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)
+    
+    # Run the API server
+    uvicorn.run(
+        "main:app",
+        host="0.0.0.0",
+        port=7860,
+        reload=False
+    )