Update app.py

app.py

@@ -1,92 +1,77 @@
-import gradio as gr
 import librosa
 import numpy as np
 import torch
 import logging
-from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2FeatureExtractor
+from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
+import gradio as gr
 
-# Initialize logging
 logging.basicConfig(level=logging.INFO)
 
-# Load the model and feature extractor
-model_path = "./"
+# Path to your Wav2Vec2 model and processor
+model_path = "./wav2vec2-sequence-classification"
 try:
     model = Wav2Vec2ForSequenceClassification.from_pretrained(model_path)
-    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_path)
-    logging.info("Model and feature extractor loaded successfully.")
+    processor = Wav2Vec2Processor.from_pretrained(model_path)
+    logging.info("Model and processor loaded successfully.")
 except Exception as e:
-    logging.error(f"Model loading failed: {e}")
+    logging.error(f"Loading model and processor failed: {e}")
     raise e
 
-def load_audio(audio_path, sr=16000):
+def preprocess_audio(file_path):
     """
-    Load an audio file and resample to the target sample rate.
+    Load and preprocess the audio file.
     """
-    try:
-        audio, _ = librosa.load(audio_path, sr=sr)
-        logging.info("Audio file loaded and resampled.")
-        return audio
-    except Exception as e:
-        logging.error(f"Failed to load audio: {e}")
-        raise e
+    # Load the audio file using librosa
+    audio, sr = librosa.load(file_path, sr=None)
+    # Resample the audio to 16 kHz (if not already at this sample rate)
+    if sr != 16000:
+        audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
+        sr = 16000
+    return audio, sr
 
-def preprocess_audio(audio):
+def audio_to_features(audio, sr):
     """
-    Preprocess the audio file to the format expected by Wav2Vec2 model.
+    Convert audio waveform to model features.
     """
-    try:
-        input_values = feature_extractor(audio, return_tensors="pt", padding="longest", sampling_rate=16000).input_values
-        logging.info("Audio file preprocessed.")
-        return input_values
-    except Exception as e:
-        logging.error(f"Audio preprocessing failed: {e}")
-        raise e
+    # Use the processor to prepare the features for the model
+    return processor(audio, sampling_rate=sr, return_tensors="pt", padding=True, truncation=True).input_values
 
-def predict(input_values):
+def classify_audio(file_path):
     """
-    Make a prediction with the Wav2Vec2 model.
+    Classify the content of the audio file.
     """
     try:
+        audio, sr = preprocess_audio(file_path)
+        input_values = audio_to_features(audio, sr)
+        
+        # Inference
         with torch.no_grad():
             logits = model(input_values).logits
-            predicted_id = torch.argmax(logits, dim=-1)
-            logging.info(f"Prediction made with id {predicted_id}")
-        return predicted_id
-    except Exception as e:
-        logging.error(f"Prediction failed: {e}")
-        raise e
-
-def get_label(prediction_id):
-    """
-    Convert the prediction ID to a meaningful label.
-    """
-    # Example of converting predicted id to label
-    # This should be adapted based on your specific model's labels
-    labels = ["label1", "label2"]  # Dummy label list for demonstration
-    try:
-        label = labels[prediction_id]
-        logging.info(f"Label obtained: {label}")
-        return label
-    except Exception as e:
-        logging.error(f"Failed to get label: {e}")
-        raise e
-
-def main(audio_file_path):
-    """
-    Load audio, preprocess, predict, and return the label.
-    """
-    try:
-        audio = load_audio(audio_file_path)
-        input_values = preprocess_audio(audio)
-        prediction_id = predict(input_values)
-        label = get_label(prediction_id)
-        return label
+        
+        # Post-processing: Convert logits to softmax to get probabilities
+        probabilities = torch.softmax(logits, dim=1).detach().numpy()
+        
+        # Assuming you have a binary classification model for simplicity
+        # Modify this part based on your actual number of classes and labels
+        labels = ['Class 0', 'Class 1']  # Example labels
+        predictions = dict(zip(labels, probabilities[0]))
+        
+        # Format the prediction output
+        prediction_output = "\n".join([f"{label}: {prob:.4f}" for label, prob in predictions.items()])
+        return prediction_output
     except Exception as e:
-        logging.error(f"Error in processing: {e}")
-        return str(e)
+        logging.error(f"Error during classification: {e}")
+        return f"Classification error: {e}"
 
-# Set up Gradio interface
-iface = gr.Interface(fn=main, inputs=gr.inputs.Audio(type="filepath"), outputs="text", title="Audio Classification")
+# Gradio interface
+iface = gr.Interface(
+    fn=classify_audio,
+    inputs=gr.inputs.Audio(source="upload", type="filepath"),
+    outputs="text",
+    title="Audio Classification with Wav2Vec2",
+    description="Upload an audio file to classify its content using a Wav2Vec2 model."
+)
 
 # Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()