fixing audio.py

tasks/audio.py

@@ -1,57 +1,52 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, HTTPException
 from datetime import datetime
-from datasets import load_dataset
+from datasets import load_dataset, get_dataset_config_names
 from sklearn.metrics import accuracy_score
 import os
 import joblib
-import librosa
 import numpy as np
-import logging
+import librosa
+from pathlib import Path
 
 from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-from dotenv import load_dotenv
-load_dotenv()
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-logger = logging.getLogger(__name__)
-
+# Router setup
 router = APIRouter()
 
 DESCRIPTION = "Chainsaw Detection Model"
 ROUTE = "/audio"
 
-# Load the trained model and scaler
-MODEL_PATH = "models/audio_model.joblib"
-model_data = joblib.load(MODEL_PATH)
-model = model_data["model"]
-scaler = model_data["scaler"]
-
-def extract_features(audio_array):
-    """Extract features from audio array."""
-    logger.debug("Extracting features from audio array...")
+# Model loading
+MODEL_PATH = Path(__file__).parent.parent / "models" / "audio_model.joblib"
+try:
+    model_data = joblib.load(MODEL_PATH)
+    model = model_data["model"]
+    scaler = model_data["scaler"]
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {e}")
+
+def extract_features(audio_array, sr=12000):
+    """Extract audio features using Librosa"""
     try:
-        # Ensure the audio is in mono
-        if len(audio_array.shape) > 1:
-            audio_array = np.mean(audio_array, axis=1)
+        # Convert to mono if stereo
+        y = np.mean(audio_array, axis=1) if len(audio_array.shape) > 1 else audio_array
         
-        # Extract MFCC features
+        # Extract MFCCs
         mfccs = librosa.feature.mfcc(
-            y=audio_array,
-            sr=12000,
+            y=y,
+            sr=sr,
             n_mfcc=13,
             n_fft=2048,
             hop_length=512
         )
         
         # Extract additional features
-        zcr = librosa.feature.zero_crossing_rate(audio_array)
-        rms = librosa.feature.rms(y=audio_array)
-        spectral_centroid = librosa.feature.spectral_centroid(y=audio_array, sr=12000)
+        zcr = librosa.feature.zero_crossing_rate(y)
+        rms = librosa.feature.rms(y=y)
+        spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr)
         
-        # Compute statistics
+        # Calculate statistics
         feature_vector = np.concatenate([
             np.mean(mfccs, axis=1),
             np.std(mfccs, axis=1),
@@ -60,95 +55,105 @@ def extract_features(audio_array):
             [np.mean(spectral_centroid)]
         ])
         
-        logger.debug("Features extracted successfully.")
         return feature_vector
+        
     except Exception as e:
-        logger.error(f"Error extracting features: {e}")
-        return None
+        raise HTTPException(status_code=400, detail=f"Feature extraction failed: {str(e)}")
 
 @router.post(ROUTE, tags=["Audio Task"], description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
-    """
-    Evaluate audio classification for rainforest sound detection.
-    
-    Current Model: Chainsaw Detection Model
-    - Uses a pre-trained RandomForestClassifier to detect chainsaw sounds.
-    """
-    logger.info("Starting audio evaluation...")
-    
-    # Get space info
-    username, space_url = get_space_info()
-    logger.info(f"Space info retrieved: username={username}, space_url={space_url}")
-
-    # Load and prepare the dataset
-    logger.info(f"Loading dataset '{request.dataset_name}'...")
     try:
-        dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
-        logger.info("Dataset loaded successfully.")
-    except Exception as e:
-        logger.error(f"Failed to load dataset: {e}")
-        raise
-
-    # Split dataset
-    logger.info(f"Splitting dataset with test_size={request.test_size}, test_seed={request.test_seed}...")
-    train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
-    test_dataset = train_test["test"]
-    logger.info(f"Dataset split into {len(test_dataset)} test samples.")
-
-    # Start tracking emissions
-    logger.info("Starting emissions tracking...")
-    tracker.start()
-    tracker.start_task("inference")
-
-    # Prepare lists to hold predictions and true labels
-    predictions = []
-    true_labels = []
-    logger.info("Starting inference on test dataset...")
-
-    # Loop through each audio sample in the test dataset
-    for i, sample in enumerate(test_dataset):
-        logger.debug(f"Processing sample {i + 1}...")
-        audio_array = sample["audio"]["array"]
-        label = sample["label"]
-
-        # Extract features
-        features = extract_features(audio_array)
-        if features is not None:
-            # Scale the features
-            features_scaled = scaler.transform([features])
-            # Make prediction
-            prediction = model.predict(features_scaled)[0]
-            predictions.append(prediction)
-            true_labels.append(label)
-        else:
-            logger.warning(f"Skipping sample {i + 1} due to feature extraction error.")
-            continue
-
-    # Stop tracking emissions
-    emissions_data = tracker.stop_task()
-    logger.info("Inference completed. Stopping emissions tracking.")
-
-    # Calculate accuracy
-    accuracy = accuracy_score(true_labels, predictions)
-    logger.info(f"Accuracy calculated: {accuracy:.4f}")
+        # Get Space info
+        username, space_url = get_space_info()
+
+        # Load dataset with proper error handling
+        try:
+            # Get available configs
+            configs = get_dataset_config_names(request.dataset_name)
+            
+            # Set up dataset loading arguments
+            dataset_args = {
+                "path": request.dataset_name,
+                "token": os.getenv("HF_TOKEN"),
+                "trust_remote_code": True
+            }
+            
+            # If configs exist, automatically use 'default' if it's the only one
+            if configs:
+                if len(configs) == 1 and configs[0] == 'default':
+                    dataset_args["name"] = "default"
+                else:
+                    raise HTTPException(
+                        status_code=400,
+                        detail=f"Config name is required for this dataset. Available configs: {configs}"
+                    )
+            
+            dataset = load_dataset(**dataset_args)
+            
+        except Exception as e:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Failed to load dataset: {str(e)}"
+            )
+
+        # Split dataset
+        split = dataset["train"].train_test_split(
+            test_size=request.test_size,
+            seed=request.test_seed
+        )
+        test_data = split["test"]
 
-    # Prepare results dictionary
-    results = {
-        "username": username,
-        "space_url": space_url,
-        "submission_timestamp": datetime.now().isoformat(),
-        "model_description": DESCRIPTION,
-        "accuracy": float(accuracy),
-        "energy_consumed_wh": emissions_data.energy_consumed * 1000,
-        "emissions_gco2eq": emissions_data.emissions * 1000,
-        "emissions_data": clean_emissions_data(emissions_data),
-        "api_route": ROUTE,
-        "dataset_config": {
-            "dataset_name": request.dataset_name,
-            "test_size": request.test_size,
-            "test_seed": request.test_seed
+        # Track emissions
+        tracker.start()
+        tracker.start_task("inference")
+        
+        # Process features
+        features = []
+        valid_samples = []
+        for sample in test_data:
+            try:
+                if 'audio' in sample and isinstance(sample['audio'], dict) and 'array' in sample['audio']:
+                    feature = extract_features(sample['audio']['array'])
+                    if feature is not None:
+                        features.append(feature)
+                        valid_samples.append(sample)
+            except Exception as e:
+                print(f"Skipping sample due to error: {e}")
+                continue
+
+        if not features:
+            raise HTTPException(
+                status_code=400,
+                detail="No valid features could be extracted from the audio samples"
+            )
+
+        # Scale features and make predictions
+        scaled_features = scaler.transform(features)
+        predictions = model.predict(scaled_features)
+        true_labels = [sample["label"] for sample in valid_samples]
+
+        # Calculate results
+        emissions_data = tracker.stop_task()
+        
+        return {
+            "username": username,
+            "space_url": space_url,
+            "submission_timestamp": datetime.now().isoformat(),
+            "model_description": DESCRIPTION,
+            "accuracy": float(accuracy_score(true_labels, predictions)),
+            "energy_consumed_wh": emissions_data.energy_consumed * 1000,
+            "emissions_gco2eq": emissions_data.emissions * 1000,
+            "emissions_data": clean_emissions_data(emissions_data),
+            "api_route": ROUTE,
+            "dataset_config": {
+                "dataset_name": request.dataset_name,
+                "test_size": request.test_size,
+                "test_seed": request.test_seed
+            }
         }
-    }
 
-    logger.info("Audio evaluation completed successfully.")
-    return results
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"An error occurred during audio evaluation: {str(e)}"
+        )