Update app.py

app.py

@@ -38,6 +38,38 @@ def convert_audio_to_wav(audio_path):
 
 
 
+import gradio as gr
+import librosa
+import numpy as np
+import whisperx
+from transformers import pipeline
+from pydub import AudioSegment
+import os
+import scipy.signal as signal
+import torch
+import pandas as pd
+from pydub.silence import detect_nonsilent
+
+hf_token = os.getenv('diarizationToken')
+
+print("Initializing Speech-to-Text Model...")
+stt_pipeline = pipeline("automatic-speech-recognition", model="boumehdi/wav2vec2-large-xlsr-moroccan-darija")
+print("Model Loaded Successfully.")
+
+# Initialize WhisperX with diarization
+device = "cuda" if torch.cuda.is_available() else "cpu"
+whisper_model = whisperx.load_model("large-v2", device)
+diarize_model = whisperx.DiarizationPipeline(use_auth_token=hf_token, device=device)
+print("WhisperX Model Loaded Successfully.")
+
+def remove_phone_tonalities(audio, sr):
+    nyquist = 0.5 * sr
+    low_cut = 300 / nyquist
+    high_cut = 3400 / nyquist
+    b, a = signal.butter(1, [low_cut, high_cut], btype='band')
+    filtered_audio = signal.filtfilt(b, a, audio)
+    return filtered_audio
+
 def process_audio(audio_path):
     print(f"Received audio file: {audio_path}")
     
@@ -71,6 +103,11 @@ def process_audio(audio_path):
         # Apply diarization (WhisperX)
         diarization = diarize_model(audio_path)
         
+        # Check if diarization is a DataFrame and process accordingly
+        if isinstance(diarization, pd.DataFrame):
+            print("Diarization is a DataFrame")
+            diarization = diarization.to_dict(orient="records")  # Convert DataFrame to a list of dicts
+        
         transcriptions = []
         for chunk in non_silent_chunks:
             chunk.export("chunk.wav", format="wav")
@@ -79,8 +116,8 @@ def process_audio(audio_path):
             
             # Match transcription segment with diarization result
             speaker_label = "Unknown"
-            for segment in diarization.itertracks(yield_label=True):
-                spk_start, spk_end, label = segment
+            for speaker in diarization:
+                spk_start, spk_end, label = speaker['start'], speaker['end'], speaker['label']
                 # Adjust timestamp matching
                 if spk_start <= (chunk.start_time / 1000) <= spk_end:  # Convert ms to seconds
                     speaker_label = label
@@ -97,7 +134,6 @@ def process_audio(audio_path):
         print(f"Error: {str(e)}")
         return f"Error: {str(e)}"
 
-
 # Create Gradio interface
 iface = gr.Interface(
     fn=process_audio,