Update audio_processing.py

audio_processing.py

@@ -1,4 +1,3 @@
-
 import whisperx
 import torch
 import numpy as np
@@ -10,50 +9,21 @@ load_dotenv()
 import logging
 import time
 from difflib import SequenceMatcher
-import spaces
-
 hf_token = os.getenv("HF_TOKEN")
 
-CHUNK_LENGTH = 5
-OVERLAP = 2
+CHUNK_LENGTH=5
+OVERLAP=0
+import whisperx
+import torch
+import numpy as np
+
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
+import spaces
 
 
-@spaces.GPU(duration=60)
-def load_whisper_model(model_size="small"):
-    logger.info(f"Loading Whisper model (size: {model_size})...")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    compute_type = "float16" if device == "cuda" else "int8"
-    try:
-        model = whisperx.load_model(model_size, device, compute_type=compute_type)
-        logger.info(f"Whisper model loaded successfully on {device}")
-        return model
-    except RuntimeError as e:
-        logger.warning(f"Failed to load Whisper model on {device}. Falling back to CPU. Error: {str(e)}")
-        device = "cpu"
-        compute_type = "int8"
-        model = whisperx.load_model(model_size, device, compute_type=compute_type)
-        logger.info("Whisper model loaded successfully on CPU")
-        return model
-
-
-@spaces.GPU(duration=60)
-def load_diarization_pipeline():
-    logger.info("Loading diarization pipeline...")
-    try:
-        pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
-        if torch.cuda.is_available():
-            pipeline = pipeline.to(torch.device("cuda"))
-        logger.info("Diarization pipeline loaded successfully")
-        return pipeline
-    except Exception as e:
-        logger.warning(f"Diarization pipeline initialization failed: {str(e)}. Diarization will not be available.")
-        return None
-
-
-def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000, overlap=OVERLAP*16000):
+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 - overlap):
         chunk = audio[i:i+chunk_size]
@@ -62,103 +32,75 @@ def preprocess_audio(audio, chunk_size=CHUNK_LENGTH*16000, overlap=OVERLAP*16000
         chunks.append(chunk)
     return chunks
 
-
-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:
-            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:
-                merged_text = merged[-1]['text'] + segment['text'][match.b + match.size:]
-                merged_translated = merged[-1].get('translated', '') + segment.get('translated', '')[match.b + match.size:]
-                
-                merged[-1]['end'] = segment['end']
-                merged[-1]['text'] = merged_text
-                if 'translated' in segment:
-                    merged[-1]['translated'] = merged_translated
-            else:
-                merged.append(segment)
-    return merged
-
-def get_most_common_speaker(diarization_result, start_time, end_time):
-    speakers = []
-    for turn, _, speaker in diarization_result.itertracks(yield_label=True):
-        if turn.start <= end_time and turn.end >= start_time:
-            speakers.append(speaker)
-    return max(set(speakers), key=speakers.count) if speakers else "Unknown"
-
-def split_audio(audio, max_duration=30):
-    sample_rate = 16000
-    max_samples = max_duration * sample_rate
-    
-    if len(audio) <= max_samples:
-        return [audio]
-    
-    splits = []
-    for i in range(0, len(audio), max_samples):
-        splits.append(audio[i:i+max_samples])
-    
-    return splits
-
-@spaces.GPU(duration=60)
-def process_audio(audio_file, translate=False, model_size="small", use_diarization=True):
-    logger.info(f"Starting audio processing: translate={translate}, model_size={model_size}, use_diarization={use_diarization}")
+@spaces.GPU()
+def process_audio(audio_file, translate=False, model_size="small"):
     start_time = time.time()
     
     try:
-        whisper_model = load_whisper_model(model_size)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {device}")
+        compute_type = "int8" if torch.cuda.is_available() else "float32"
         audio = whisperx.load_audio(audio_file)
-        audio_splits = split_audio(audio)
+        model = whisperx.load_model(model_size, device, compute_type=compute_type)
+
+        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})
 
-        diarization_result = None
-        if use_diarization:
-            diarization_pipeline = load_diarization_pipeline()
-            if diarization_pipeline is not None:
-                try:
-                    diarization_result = diarization_pipeline({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000})
-                except Exception as e:
-                    logger.warning(f"Diarization failed: {str(e)}. Proceeding without diarization.")
+        chunks = preprocess_audio(audio)
 
         language_segments = []
         final_segments = []
         
-        for i, audio_split in enumerate(audio_splits):
-            logger.info(f"Processing split {i+1}/{len(audio_splits)}")
+        overlap_duration = OVERLAP  # 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
             
-            result = whisper_model.transcribe(audio_split)
-            lang = result["language"]
+                # 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
             
-            for segment in result["segments"]:
-                segment_start = segment["start"] + (i * 30)
-                segment_end = segment["end"] + (i * 30)
-                
-                speaker = "Unknown"
-                if diarization_result is not None:
-                    speaker = get_most_common_speaker(diarization_result, segment_start, segment_end)
-                
-                final_segment = {
+                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": speaker,
-                    "text": segment["text"],
+                    "speaker": max(set(speakers), key=speakers.count) if speakers else "Unknown",
+                    "text": t_seg["text"],
                 }
                 
                 if translate:
-                    translation = whisper_model.transcribe(audio_split[int(segment["start"]*16000):int(segment["end"]*16000)], task="translate")
-                    final_segment["translated"] = translation["text"]
-                
-                final_segments.append(final_segment)
+                    segment["translated"] = result_translate["segments"][j]["text"]
             
+                final_segments.append(segment)
+
             language_segments.append({
                 "language": lang,
-                "start": i * 30,
-                "end": min((i + 1) * 30, len(audio) / 16000)
+                "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")
 
         final_segments.sort(key=lambda x: x["start"])
         merged_segments = merge_nearby_segments(final_segments)
@@ -166,7 +108,38 @@ def process_audio(audio_file, translate=False, model_size="small", use_diarizati
         end_time = time.time()
         logger.info(f"Total processing time: {end_time - start_time:.2f} seconds")
 
-        return language_segments, merged_segments
+        return language_segments, final_segments
     except Exception as e:
         logger.error(f"An error occurred during audio processing: {str(e)}")
-        raise
+        raise
+
+def merge_nearby_segments(segments, time_threshold=0.5, similarity_threshold=0.9):
+    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] ({segment['language']}) {segment['speaker']}:")
+        print(f"Original: {segment['text']}")
+        if 'translated' in segment:
+            print(f"Translated: {segment['translated']}")
+        print()