Update app.py

app.py

@@ -2,10 +2,10 @@ import gradio as gr
 from pyannote.audio import Pipeline
 import torch
 import os
-import numpy as np
-from pydub import AudioSegment
-import io
 import zipfile
+import tempfile
+import shutil
+import librosa
 
 hf_token = os.getenv("HF_TOKEN")
 
@@ -13,59 +13,61 @@ hf_token = os.getenv("HF_TOKEN")
 pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
 pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
-def diarize_and_split(audio, sr):
-    # Convert to mono if stereo
-    if len(audio.shape) > 1:
-        audio = np.mean(audio, axis=1)
-    
-    # Perform diarization
-    diarization = pipeline({"waveform": torch.from_numpy(audio), "sample_rate": sr})
-    
-    # Create an AudioSegment from the numpy array
-    audio_segment = AudioSegment(
-        audio.tobytes(), 
-        frame_rate=sr,
-        sample_width=audio.dtype.itemsize, 
-        channels=1
-    )
-    
-    speaker_segments = {}
-    
-    for turn, _, speaker in diarization.itertracks(yield_label=True):
-        start_ms = int(turn.start * 1000)
-        end_ms = int(turn.end * 1000)
-        segment = audio_segment[start_ms:end_ms]
+def process_zip(zip_file):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Step 1: Extract the zip file
+        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+            zip_ref.extractall(temp_dir)
         
-        if speaker not in speaker_segments:
-            speaker_segments[speaker] = []
-        speaker_segments[speaker].append(segment)
-    
-    # Create zip files for each speaker
-    zip_files = {}
-    for speaker, segments in speaker_segments.items():
-        zip_buffer = io.BytesIO()
-        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
-            for i, segment in enumerate(segments):
-                segment_buffer = io.BytesIO()
-                segment.export(segment_buffer, format="wav")
-                zip_file.writestr(f"{speaker}_segment_{i}.wav", segment_buffer.getvalue())
+        # Create directories for each speaker
+        speaker1_dir = os.path.join(temp_dir, "speaker1")
+        speaker2_dir = os.path.join(temp_dir, "speaker2")
+        os.makedirs(speaker1_dir, exist_ok=True)
+        os.makedirs(speaker2_dir, exist_ok=True)
         
-        zip_buffer.seek(0)
-        zip_files[f"{speaker}.zip"] = zip_buffer.getvalue()
-    
-    return zip_files
-
-def process_audio(audio):
-    sr, audio_data = audio
-    zip_files = diarize_and_split(audio_data, sr)
-    return list(zip_files.values())
+        # Step 2: Analyze each audio file
+        for filename in os.listdir(temp_dir):
+            if filename.endswith(('.wav', '.mp3', '.ogg', '.flac')):
+                file_path = os.path.join(temp_dir, filename)
+                
+                # Load audio file
+                waveform, sample_rate = librosa.load(file_path, sr=None)
+                
+                # Perform diarization
+                diarization = pipeline({"waveform": torch.from_numpy(waveform), "sample_rate": sample_rate})
+                
+                # Determine dominant speaker
+                speaker_times = {1: 0, 2: 0}
+                for turn, _, speaker in diarization.itertracks(yield_label=True):
+                    speaker_num = int(speaker.split('_')[-1])
+                    speaker_times[speaker_num] += turn.end - turn.start
+                
+                dominant_speaker = 1 if speaker_times[1] > speaker_times[2] else 2
+                
+                # Move file to appropriate speaker directory
+                if dominant_speaker == 1:
+                    shutil.move(file_path, os.path.join(speaker1_dir, filename))
+                else:
+                    shutil.move(file_path, os.path.join(speaker2_dir, filename))
+        
+        # Step 3: Create zip files for each speaker
+        speaker1_zip = os.path.join(temp_dir, "speaker1.zip")
+        speaker2_zip = os.path.join(temp_dir, "speaker2.zip")
+        
+        shutil.make_archive(os.path.join(temp_dir, "speaker1"), 'zip', speaker1_dir)
+        shutil.make_archive(os.path.join(temp_dir, "speaker2"), 'zip', speaker2_dir)
+        
+        return speaker1_zip, speaker2_zip
 
 iface = gr.Interface(
-    fn=process_audio,
-    inputs=gr.Audio(type="numpy"),
-    outputs=[gr.File(label="Speaker Zip Files") for _ in range(10)],  # Assuming max 10 speakers
-    title="Speaker Diarization and Audio Splitting",
-    description="Upload an audio file to split it into separate files for each speaker."
+    fn=process_zip,
+    inputs=gr.File(label="Upload ZIP file containing audio files"),
+    outputs=[
+        gr.File(label="Speaker 1 Audio Files"),
+        gr.File(label="Speaker 2 Audio Files")
+    ],
+    title="Speaker Diarization and Audio Sorting",
+    description="Upload a ZIP file containing audio files. The system will analyze each file and sort them into two groups based on the dominant speaker."
 )
 
 iface.launch()