Update audio_processing.py

audio_processing.py

@@ -1,112 +1,105 @@
+import whisperx
 import torch
-import whisper
-import torchaudio as ta
-import gradio as gr
-from model_utils import get_processor, get_model, get_whisper_model_small, get_device
-from config import SAMPLING_RATE, CHUNK_LENGTH_S
-import spaces
-
-
-@spaces.GPU
-def load_and_resample_audio(audio):
-    if isinstance(audio, str):  # If audio is a file path
-        waveform, sample_rate = ta.load(audio)
-    else:  # If audio is already loaded (sample_rate, waveform)
-        sample_rate, waveform = audio
-        waveform = torch.tensor(waveform).float()
-
-    if sample_rate != SAMPLING_RATE:
-        waveform = ta.functional.resample(waveform, sample_rate, SAMPLING_RATE)
-
-    # Ensure the audio is in the correct shape (mono)
-    if waveform.dim() > 1 and waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0, keepdim=True)
-    elif waveform.dim() == 1:
-        waveform = waveform.unsqueeze(0)
-
-    return waveform, SAMPLING_RATE
-
-
-@spaces.GPU
-def detect_language(waveform):
-    whisper_model = get_whisper_model_small()
-
-    # Use Whisper's preprocessing
-    audio_tensor = whisper.pad_or_trim(waveform.squeeze())
-    mel = whisper.log_mel_spectrogram(audio_tensor).to(whisper_model.device)
-
-    # Detect language
-    _, probs = whisper_model.detect_language(mel)
-    detected_lang = max(probs, key=probs.get)
-
-    print(f"Audio shape: {audio_tensor.shape}")
-    print(f"Mel spectrogram shape: {mel.shape}")
-    print(f"Detected language: {detected_lang}")
-    print("Language probabilities:", probs)
-
-    return detected_lang
-
-
-@spaces.GPU
-def process_long_audio(waveform, sample_rate, task="transcribe", language=None):
-    input_length = waveform.shape[1]
-    chunk_length = int(CHUNK_LENGTH_S * sample_rate)
+import numpy as np
+from scipy.signal import resample
+import numpy as np
+import whisperx
+from pyannote.audio import Pipeline
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+hf_token = os.getenv("HF_TOKEN")
+import whisperx
+import torch
+import numpy as np
 
-    chunks = [waveform[:, i:i + chunk_length] for i in range(0, input_length, chunk_length)]
+import whisperx
+import torch
+import numpy as np
 
-    processor = get_processor()
-    model = get_model()
-    device = get_device()
+import whisperx
+import torch
+import numpy as np
+CHUNK_LENGTH= 30
 
-    results = []
-    for chunk in chunks:
-        input_features = processor(chunk.squeeze(), sampling_rate=sample_rate, return_tensors="pt").input_features.to(
-            device)
 
-        with torch.no_grad():
-            if task == "translate":
-                forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task="translate")
-                generated_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
+import whisperx
+import torch
+import numpy as np
+
+def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000):  # 30 seconds at 16kHz
+    chunks = []
+    for i in range(0, len(audio), chunk_size):
+        chunk = audio[i:i+chunk_size]
+        if len(chunk) < chunk_size:
+            chunk = np.pad(chunk, (0, chunk_size - len(chunk)))
+        chunks.append(chunk)
+    return chunks
+
+def process_audio(audio_file):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    compute_type = "float32"
+    audio = whisperx.load_audio(audio_file)
+    model = whisperx.load_model("small", device, compute_type=compute_type)
+
+    # Initialize speaker diarization pipeline
+    diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
+    diarization_pipeline = diarization_pipeline.to(torch.device(device))
+
+    # Perform diarization on the entire audio
+    diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
+
+
+    # Preprocess audio into consistent chunks
+    chunks = preprocess_audio(audio)
+
+    language_segments = []
+    final_segments = []
+    
+    for i, chunk in enumerate(chunks):
+        # Detect language for this chunk
+        lang = model.detect_language(chunk)
+        
+        # Transcribe this chunk
+        result = model.transcribe(chunk, language=lang)
+        
+        chunk_start_time = i * 5  # Each chunk is 30 seconds
+        
+        # Adjust timestamps and add language information
+        for segment in result["segments"]:
+            segment_start = chunk_start_time + segment["start"]
+            segment_end = chunk_start_time + segment["end"]
+            segment["start"] = segment_start
+            segment["end"] = segment_end
+            segment["language"] = lang
+            
+            speakers = []
+            for turn, track, speaker in diarization_result.itertracks(yield_label=True):
+                if turn.start <= segment_end and turn.end >= segment_start:
+                    speakers.append(speaker)
+            if speakers:
+                segment["speaker"] = max(set(speakers), key=speakers.count)
             else:
-                generated_ids = model.generate(input_features)
-
-        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)
-        results.extend(transcription)
-
-        # Clear GPU cache
-        torch.cuda.empty_cache()
-
-    return " ".join(results)
-
-
-@spaces.GPU
-def process_audio(audio):
-    if audio is None:
-        return "No file uploaded", "", ""
-
-    waveform, sample_rate = load_and_resample_audio(audio)
-
-    detected_lang = detect_language(waveform)
-    transcription = process_long_audio(waveform, sample_rate, task="transcribe")
-    translation = process_long_audio(waveform, sample_rate, task="translate", language=detected_lang)
-
-    return detected_lang, transcription, translation
-
-
-# Gradio interface
-iface = gr.Interface(
-    fn=process_audio,
-    inputs=gr.Audio(),
-    outputs=[
-        gr.Textbox(label="Detected Language"),
-        gr.Textbox(label="Transcription", lines=5),
-        gr.Textbox(label="Translation", lines=5)
-    ],
-    title="Audio Transcription and Translation",
-    description="Upload an audio file to detect its language, transcribe, and translate it.",
-    allow_flagging="never",
-    css=".output-textbox { font-family: 'Noto Sans Devanagari', sans-serif; font-size: 18px; }"
-)
-
-if __name__ == "__main__":
-    iface.launch()
+                segment["speaker"] = "Unknown"
+
+            final_segments.append(segment)
+        # Add language segment
+        language_segments.append({
+            "language": lang,
+            "start": chunk_start_time,
+            "end": chunk_start_time + 5
+        })
+
+    return language_segments, final_segments
+
+def print_results(language, language_probs, segments):
+    print(f"Detected Language: {language}")
+    print("Language Probabilities:")
+    for lang, prob in language_probs.items():
+        print(f"  {lang}: {prob:.4f}")
+    
+    print("\nTranscription:")
+    for segment in segments:
+        print(f"[{segment['start']:.2f}s - {segment['end']:.2f}s] Speaker {segment['speaker']}: {segment['text']}")