added translation, and model selection

app.py

@@ -1,24 +1,32 @@
 import gradio as gr
 from audio_processing import process_audio, print_results
-def transcribe_audio(audio_file):
-    language_segments, final_segments = process_audio(audio_file)
+
+def transcribe_audio(audio_file, translate, model_size):
+    language_segments, final_segments = process_audio(audio_file, translate=translate, model_size=model_size)
     
     output = "Detected language changes:\n\n"
     for segment in language_segments:
         output += f"Language: {segment['language']}\n"
         output += f"Time: {segment['start']:.2f}s - {segment['end']:.2f}s\n\n"
 
-    output += "Transcription with language detection and speaker diarization:\n\n"
+    output += f"Transcription with language detection and speaker diarization (using {model_size} model):\n\n"
     for segment in final_segments:
-        output += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}) Speaker {segment['speaker']}: {segment['text']}\n"
-        # output += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}): {segment['text']}\n"
+        output += f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}) {segment['speaker']}:\n"
+        output += f"Original: {segment['text']}\n"
+        if translate:
+            output += f"Translated: {segment['translated']}\n"
+        output += "\n"
     return output
 
 iface = gr.Interface(
     fn=transcribe_audio,
-    inputs=gr.Audio(type="filepath"),
+    inputs=[
+        gr.Audio(type="filepath"),
+        gr.Checkbox(label="Enable Translation"),
+        gr.Dropdown(choices=["tiny", "base", "small", "medium", "large","large-v2","large-v3"], label="Whisper Model Size", value="small")
+    ],
     outputs="text",
-    title="WhisperX Audio Transcription"
+    title="WhisperX Audio Transcription and Translation"
 )
 
 iface.launch()

audio_processing.py

@@ -2,168 +2,139 @@ 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()
-
+import logging
+import time
+from difflib import SequenceMatcher
 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
-
+OVERLAP=2
 import whisperx
 import torch
 import numpy as np
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
-def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000):  # 30 seconds at 16kHz
+def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000, overlap=OVERLAP*16000):  # 2 seconds overlap
     chunks = []
-    for i in range(0, len(audio), chunk_size):
+    for i in range(0, len(audio), chunk_size - overlap):
         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))
+def process_audio(audio_file, translate=False, model_size="small"):
+    start_time = time.time()
+    
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        compute_type = "float32"
+        audio = whisperx.load_audio(audio_file)
+        model = whisperx.load_model(model_size, device, compute_type=compute_type)
 
-    # Perform diarization on the entire audio
-    diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
+        diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
+        diarization_pipeline = diarization_pipeline.to(torch.device(device))
 
+        diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
 
-    # Preprocess audio into consistent chunks
-    chunks = preprocess_audio(audio)
+        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)
+        language_segments = []
+        final_segments = []
         
-        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
+        overlap_duration = 2  # 2 seconds overlap
+        for i, chunk in enumerate(chunks):
+            chunk_start_time = i * (CHUNK_LENGTH - overlap_duration)
+            chunk_end_time = chunk_start_time + CHUNK_LENGTH
+            logger.info(f"Processing chunk {i+1}/{len(chunks)}")
+            lang = model.detect_language(chunk)
+            result_transcribe = model.transcribe(chunk, language=lang)
+            if translate:
+                result_translate = model.transcribe(chunk, task="translate")
+            chunk_start_time = i * (CHUNK_LENGTH - overlap_duration)
+            for j, t_seg in enumerate(result_transcribe["segments"]):
+                segment_start = chunk_start_time + t_seg["start"]
+                segment_end = chunk_start_time + t_seg["end"]
+                # Skip segments in the overlapping region of the previous chunk
+                if i > 0 and segment_end <= chunk_start_time + overlap_duration:
+                    print(f"Skipping segment in overlap with previous chunk: {segment_start:.2f} - {segment_end:.2f}")
+                    continue
             
-            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"
+                # Skip segments in the overlapping region of the next chunk
+                if i < len(chunks) - 1 and segment_start >= chunk_end_time - overlap_duration:
+                    print(f"Skipping segment in overlap with next chunk: {segment_start:.2f} - {segment_end:.2f}")
+                    continue
+            
+                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)
+            
+                segment = {
+                    "start": segment_start,
+                    "end": segment_end,
+                    "language": lang,
+                    "speaker": max(set(speakers), key=speakers.count) if speakers else "Unknown",
+                    "text": t_seg["text"],
+                }
+                
+                if translate:
+                    segment["translated"] = result_translate["segments"][j]["text"]
+            
+                final_segments.append(segment)
 
-            final_segments.append(segment)
-        # Add language segment
-        language_segments.append({
-            "language": lang,
-            "start": chunk_start_time,
-            "end": chunk_start_time + 5
-        })
+            language_segments.append({
+                "language": lang,
+                "start": chunk_start_time,
+                "end": chunk_start_time + CHUNK_LENGTH
+            })
+            chunk_end_time = time.time()
+            logger.info(f"Chunk {i+1} processed in {chunk_end_time - chunk_start_time:.2f} seconds")
 
-    return language_segments, final_segments
+        final_segments.sort(key=lambda x: x["start"])
+        merged_segments = merge_nearby_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:")
+        end_time = time.time()
+        logger.info(f"Total processing time: {end_time - start_time:.2f} seconds")
+
+        return language_segments, merged_segments
+    except Exception as e:
+        logger.error(f"An error occurred during audio processing: {str(e)}")
+        raise
+
+def merge_nearby_segments(segments, time_threshold=0.5, similarity_threshold=0.7):
+    merged = []
+    for segment in segments:
+        if not merged or segment['start'] - merged[-1]['end'] > time_threshold:
+            merged.append(segment)
+        else:
+            # Find the overlap
+            matcher = SequenceMatcher(None, merged[-1]['text'], segment['text'])
+            match = matcher.find_longest_match(0, len(merged[-1]['text']), 0, len(segment['text']))
+            
+            if match.size / len(segment['text']) > similarity_threshold:
+                # Merge the segments
+                merged_text = merged[-1]['text'] + segment['text'][match.b + match.size:]
+                merged_translated = merged[-1]['translated'] + segment['translated'][match.b + match.size:]
+                
+                merged[-1]['end'] = segment['end']
+                merged[-1]['text'] = merged_text
+                merged[-1]['translated'] = merged_translated
+            else:
+                # If no significant overlap, append as a new segment
+                merged.append(segment)
+    return merged
+
+def print_results(segments):
     for segment in segments:
-        print(f"[{segment['start']:.2f}s - {segment['end']:.2f}s] Speaker {segment['speaker']}: {segment['text']}")
+        print(f"[{segment['start']:.2f}s - {segment['end']:.2f}s] ({segment['language']}) {segment['speaker']}:")
+        print(f"Original: {segment['text']}")
+        if 'translated' in segment:
+            print(f"Translated: {segment['translated']}")
+        print()

requirements.txt

@@ -12,4 +12,6 @@ torchaudio>=2
 faster-whisper==1.0.0
 setuptools>=65
 nltk
-python-dotenv
+python-dotenv
+difflib
+pydub