Update audio_feature_extraction.py

audio_feature_extraction.py

@@ -1,41 +1,23 @@
-import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
-from audio_feature_extraction import extract_features
+import torch
+import librosa
+from transformers import Wav2Vec2Processor, Wav2Vec2Model
 
-def plot_spectrum(audio_path):
-    # Extract features from the audio file
-    mean_features = extract_features(audio_path)
+def extract_features(audio_path):
+    # Load the audio file
+    audio_data, sr = librosa.load(audio_path, sr=16000)  # Ensure 16 kHz sampling rate
 
-    # Perform FFT to obtain frequency components
-    fft_spectrum = np.fft.fft(mean_features)
-    frequencies = np.fft.fftfreq(len(fft_spectrum), d=1/16000)
-    magnitude_spectrum = np.abs(fft_spectrum)
+    # Load Wav2Vec2 processor and model
+    processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+    model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")
 
-    # Plot the frequency spectrum
-    plt.figure(figsize=(12, 6))
-    plt.plot(frequencies[:len(frequencies)//2], magnitude_spectrum[:len(magnitude_spectrum)//2])
-    plt.xlabel("Frequency (Hz)")
-    plt.ylabel("Magnitude")
-    plt.title("Frequency Spectrum of the Audio File")
-    plt.grid()
-    plt.tight_layout()
+    # Process the audio data
+    input_values = processor(audio_data, return_tensors="pt", sampling_rate=16000).input_values
 
-    # Save the plot to a file
-    plot_path = "spectrum_plot.png"
-    plt.savefig(plot_path)
-    plt.close()
+    # Extract features
+    with torch.no_grad():
+        features = model(input_values).last_hidden_state
 
-    return plot_path
+    # Compute the mean feature vector
+    mean_features = features.mean(dim=1).squeeze().cpu().numpy()
 
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=plot_spectrum,
-    inputs=gr.Audio(source="upload", type="filepath"),
-    outputs=gr.Image(type="filepath"),
-    title="Audio Feature Extraction with Wav2Vec2",
-    description="Upload an audio file to extract features and view the frequency spectrum."
-)
-
-if __name__ == "__main__":
-    iface.launch()
+    return mean_features