Update app.py

app.py

@@ -3,61 +3,43 @@ import librosa
 import numpy as np
 import torch
 import matplotlib.pyplot as plt
-from transformers import AutoModelForAudioClassification, ASTFeatureExtractor, Wav2Vec2Processor, Wav2Vec2ForCTC
+from transformers import AutoModelForAudioClassification, ASTFeatureExtractor
 import random
 import tempfile
 import logging
+import os
 
 logging.basicConfig(level=logging.DEBUG, filename='app.log', filemode='w', format='%(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Load Wav2Vec 2.0 models
-wav2vec_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
-wav2vec_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
-
-# Original model and feature extractor loading
 model = AutoModelForAudioClassification.from_pretrained("./")
 feature_extractor = ASTFeatureExtractor.from_pretrained("./")
 
 def plot_waveform(waveform, sr):
-    try:
-        plt.figure(figsize=(12, 4))
-        plt.title('Waveform')
-        plt.ylabel('Amplitude')
-        plt.plot(np.linspace(0, len(waveform) / sr, len(waveform)), waveform)
-        plt.xlabel('Time (s)')
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir='./')
-        plt.savefig(temp_file.name)
-        plt.close()
-        
-        file_size = os.path.getsize(temp_file.name)
-        logger.debug(f"Waveform image generated: {temp_file.name}, Size: {file_size} bytes")
-        
-        return temp_file.name
-    except Exception as e:
-        logger.error(f"Error generating waveform image: {e}")
-        raise
+    plt.figure(figsize=(24, 8))  # Doubled size for larger visuals
+    plt.title('Waveform')
+    plt.ylabel('Amplitude')
+    plt.plot(np.linspace(0, len(waveform) / sr, len(waveform)), waveform)
+    plt.xlabel('Time (s)')
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir='./')
+    plt.savefig(temp_file.name)
+    plt.close()
+    logger.debug(f"Waveform image generated: {temp_file.name}, Size: {os.path.getsize(temp_file.name)} bytes")
+    return temp_file.name
 
 def plot_spectrogram(waveform, sr):
-    try:
-        S = librosa.feature.melspectrogram(y=waveform, sr=sr, n_mels=128)
-        S_DB = librosa.power_to_db(S, ref=np.max)
-        plt.figure(figsize=(12, 6))
-        librosa.display.specshow(S_DB, sr=sr, x_axis='time', y_axis='mel')
-        plt.title('Mel Spectrogram')
-        plt.colorbar(format='%+2.0f dB')
-        plt.tight_layout()
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir='./')
-        plt.savefig(temp_file.name)
-        plt.close()
-        
-        file_size = os.path.getsize(temp_file.name)
-        logger.debug(f"Spectrogram image generated: {temp_file.name}, Size: {file_size} bytes")
-        
-        return temp_file.name
-    except Exception as e:
-        logger.error(f"Error generating spectrogram image: {e}")
-        raise
+    S = librosa.feature.melspectrogram(y=waveform, sr=sr, n_mels=128)
+    S_DB = librosa.power_to_db(S, ref=np.max)
+    plt.figure(figsize=(24, 12))  # Doubled size for larger visuals
+    librosa.display.specshow(S_DB, sr=sr, x_axis='time', y_axis='mel')
+    plt.title('Mel Spectrogram')
+    plt.colorbar(format='%+2.0f dB')
+    plt.tight_layout()
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir='./')
+    plt.savefig(temp_file.name)
+    plt.close()
+    logger.debug(f"Spectrogram image generated: {temp_file.name}, Size: {os.path.getsize(temp_file.name)} bytes")
+    return temp_file.name
 
 def custom_feature_extraction(audio, sr=16000, target_length=1024):
     features = feature_extractor(audio, sampling_rate=sr, return_tensors="pt", padding="max_length", max_length=target_length)
@@ -67,72 +49,37 @@ def apply_time_shift(waveform, max_shift_fraction=0.1):
     shift = random.randint(-int(max_shift_fraction * len(waveform)), int(max_shift_fraction * len(waveform)))
     return np.roll(waveform, shift)
 
-def transcribe_audio(audio_file_path):
-    waveform, _ = librosa.load(audio_file_path, sr=wav2vec_processor.feature_extractor.sampling_rate, mono=True)
-    input_values = wav2vec_processor(waveform, return_tensors="pt", padding="longest").input_values
-    with torch.no_grad():
-        logits = wav2vec_model(input_values).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = wav2vec_processor.batch_decode(predicted_ids)
-    return transcription
-
 def predict_voice(audio_file_path):
-    try:
-        transcription = transcribe_audio(audio_file_path)
-        
-        waveform, sample_rate = librosa.load(audio_file_path, sr=feature_extractor.sampling_rate, mono=True)
-        augmented_waveform = apply_time_shift(waveform)
-        
-        original_features = custom_feature_extraction(waveform, sr=sample_rate)
-        augmented_features = custom_feature_extraction(augmented_waveform, sr=sample_rate)
-        
-        with torch.no_grad():
-            outputs_original = model(original_features)
-            outputs_augmented = model(augmented_features)
-        
-        logits = (outputs_original.logits + outputs_augmented.logits) / 2
-        predicted_index = logits.argmax()
-        original_label = model.config.id2label[predicted_index.item()]
-        confidence = torch.softmax(logits, dim=1).max().item() * 100
-        
-        label_mapping = {
-            "Spoof": "AI-generated Clone",
-            "Bonafide": "Real Human Voice"
-        }
-        new_label = label_mapping.get(original_label, "Unknown")
-        
-        waveform_plot = plot_waveform(waveform, sample_rate)
-        spectrogram_plot = plot_spectrogram(waveform, sample_rate)
-        
-        return (
-            f"The voice is classified as '{new_label}' with a confidence of {confidence:.2f}%.",
+    waveform, _ = librosa.load(audio_file_path, sr=16000, mono=True)  # Ensure all audio is resampled to 16kHz
+    augmented_waveform = apply_time_shift(waveform)
+    original_features = custom_feature_extraction(waveform, sr=16000)  # Adjusted sample rate to 16kHz
+    augmented_features = custom_feature_extraction(augmented_waveform, sr=16000)  # Adjusted sample rate to 16kHz
+    with torch.no_grad():
+        outputs_original = model(original_features)
+        outputs_augmented = model(augmented_features)
+    logits = (outputs_original.logits + outputs_augmented.logits) / 2
+    predicted_index = logits.argmax()
+    original_label = model.config.id2label[predicted_index.item()]
+    confidence = torch.softmax(logits, dim=1).max().item() * 100
+    label_mapping = {"Spoof": "AI-generated Clone", "Bonafide": "Real Human Voice"}
+    new_label = label_mapping.get(original_label, "Unknown")
+    waveform_plot = plot_waveform(waveform, 16000)  # Adjusted sample rate to 16kHz
+    spectrogram_plot = plot_spectrogram(waveform, 16000)  # Adjusted sample rate to 16kHz
+    return (f"The voice is classified as '{new_label}' with a confidence of {confidence:.2f}%.",
             waveform_plot,
-            spectrogram_plot,
-            transcription[0]  # Assuming transcription returns a list with a single string
-        )
-     except Exception as e:
-        logger.error(f"Error during voice prediction: {e}")
-        return f"Error during processing: {e}", None, None, ""
+            spectrogram_plot)
 
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown("## Voice Clone Detection")
     gr.Markdown("Detects whether a voice is real or an AI-generated clone. Upload an audio file to see the results.")
-    
     with gr.Row():
         audio_input = gr.Audio(label="Upload Audio File", type="filepath")
-
+    detect_button = gr.Button("Detect Voice Clone")
     with gr.Row():
         prediction_output = gr.Textbox(label="Prediction")
-        transcription_output = gr.Textbox(label="Transcription")  # Fixed indentation
+    with gr.Row():
         waveform_output = gr.Image(label="Waveform")
         spectrogram_output = gr.Image(label="Spectrogram")
+    detect_button.click(fn=predict_voice, inputs=[audio_input], outputs=[prediction_output, waveform_output, spectrogram_output])
 
-    detect_button = gr.Button("Detect Voice Clone")
-    detect_button.click(
-        fn=predict_voice,
-        inputs=[audio_input],
-        outputs=[prediction_output, waveform_output, spectrogram_output, transcription_output]
-    )
-
-# Launch the interface
-demo.launch()
+demo.launch()