asr_timestamp_transcription+diarization

audio_processing.py

@@ -0,0 +1,169 @@
+import whisperx
+import torch
+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
+
+import whisperx
+import torch
+import numpy as np
+
+import whisperx
+import torch
+import numpy as np
+CHUNK_LENGTH=5
+
+# 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)
+
+#     # Initial transcription
+#     result = model.transcribe(audio, batch_size=8)
+
+#     # Sliding window for language detection
+#     window_size = 5  # seconds
+#     step_size = 1  # seconds
+#     sample_rate = 16000
+
+#     language_probs = []
+#     audio_duration = len(audio) / sample_rate
+
+#     if audio_duration <= window_size:
+#         # If audio is shorter than or equal to window size, detect language for entire audio
+#         lang = model.detect_language(audio)
+#         language_probs.append((0, lang))
+#     else:
+#         for i in range(0, len(audio) - window_size * sample_rate + 1, step_size * sample_rate):
+#             window = audio[i:i + window_size * sample_rate]
+#             lang = model.detect_language(window)
+#             language_probs.append((i / sample_rate, lang))
+
+#     # Detect language changes
+#     language_segments = []
+#     current_lang = language_probs[0][1]
+#     start_time = 0
+#     for time, lang in language_probs[1:]:
+#         if lang != current_lang:
+#             language_segments.append({
+#                 "language": current_lang,
+#                 "start": start_time,
+#                 "end": time
+#             })
+#             current_lang = lang
+#             start_time = time
+
+#     # Add the last segment
+#     language_segments.append({
+#         "language": current_lang,
+#         "start": start_time,
+#         "end": audio_duration
+#     })
+
+#     # Re-transcribe each language segment
+#     final_segments = []
+#     for segment in language_segments:
+#         start_sample = int(segment["start"] * sample_rate)
+#         end_sample = int(segment["end"] * sample_rate)
+#         segment_audio = audio[start_sample:end_sample]
+        
+#         segment_result = model.transcribe(segment_audio, language=segment["language"])
+        
+#         for seg in segment_result["segments"]:
+#             seg["start"] += segment["start"]
+#             seg["end"] += segment["start"]
+#             seg["language"] = segment["language"]
+#             final_segments.append(seg)
+
+#     return language_segments, final_segments
+
+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:
+                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']}")