Update app.py

app.py

@@ -1,86 +1,81 @@
-import gradio as gr
-from pyannote.audio import Pipeline
-import torch
 import os
 import zipfile
-import tempfile
 import shutil
-from pydub import AudioSegment
-import numpy as np
+import torch
+import torchaudio
+from pyannote.audio import Pipeline
+from pyannote.core import Segment
+import gradio as gr
+
+# Load the pre-trained model using your Hugging Face access token
+HUGGINGFACE_ACCESS_TOKEN = 'YOUR_ACCESS_TOKEN'
+pipeline = Pipeline.from_pretrained(
+    "pyannote/speaker-diarization-3.1",
+    use_auth_token=HUGGINGFACE_ACCESS_TOKEN)
+
+# Function to unzip the uploaded file
+def unzip_files(zip_fp, extract_to):
+    with zipfile.ZipFile(zip_fp, 'r') as z:
+        z.extractall(extract_to)
+
+# Function to zip files to a zip file
+def zip_files(input_dir, zip_fp):
+    with zipfile.ZipFile(zip_fp, 'w') as z:
+        for folder_name, subfolders, filenames in os.walk(input_dir):
+            for filename in filenames:
+                file_path = os.path.join(folder_name, filename)
+                z.write(file_path, os.path.relpath(file_path, input_dir))
+
+# Function to classify and group files by speaker
+def classify_and_group_speakers(zip_file):
+    # Step 1: Create temporary directories
+    extract_dir = 'extract_temp'
+    speaker1_dir = 'speaker1_temp'
+    speaker2_dir = 'speaker2_temp'
+    os.makedirs(extract_dir, exist_ok=True)
+    os.makedirs(speaker1_dir, exist_ok=True)
+    os.makedirs(speaker2_dir, exist_ok=True)
+
+    # Step 2: Extract uploaded zip file
+    unzip_files(zip_file.name, extract_dir)
+
+    # Step 3: Analyze each audio file and determine the speaker
+    for audio_file in os.listdir(extract_dir):
+        audio_fp = os.path.join(extract_dir, audio_file)
+        waveform, sample_rate = torchaudio.load(audio_fp)
+        diarization = pipeline({"waveform": waveform, "sample_rate": sample_rate})
+
+        # Check which speaker is dominant
+        speaker1_segments = [segment for segment, _, label in diarization.itertracks(yield_label=True) if label == 'SPEAKER_00']
+        speaker2_segments = [segment for segment, _, label in diarization.itertracks(yield_label=True) if label == 'SPEAKER_01']
+
+        speaker1_duration = sum([segment.duration for segment in speaker1_segments])
+        speaker2_duration = sum([segment.duration for segment in speaker2_segments])
+
+        if speaker1_duration > speaker2_duration:
+            shutil.copy(audio_fp, speaker1_dir)
+        else:
+            shutil.copy(audio_fp, speaker2_dir)
 
-# Monkey patch numpy to add NAN attribute
-np.NAN = np.nan
+    # Step 4: Zip the grouped files
+    speaker1_zip = 'speaker1.zip'
+    speaker2_zip = 'speaker2.zip'
+    zip_files(speaker1_dir, speaker1_zip)
+    zip_files(speaker2_dir, speaker2_zip)
 
-hf_token = os.getenv("HF_TOKEN")
+    # Step 5: Clean up temporary directories
+    shutil.rmtree(extract_dir)
+    shutil.rmtree(speaker1_dir)
+    shutil.rmtree(speaker2_dir)
 
-# Initialize the diarization pipeline
-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"))
+    return speaker1_zip, speaker2_zip
 
-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)
-        
-        # 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)
-        
-        # Step 2: Analyze each audio file
-        for filename in os.listdir(temp_dir):
-            if filename.lower().endswith(('.wav', '.mp3', '.ogg', '.flac')):
-                file_path = os.path.join(temp_dir, filename)
-                
-                # Load audio file
-                audio = AudioSegment.from_file(file_path)
-                samples = np.array(audio.get_array_of_samples())
-                
-                # Convert to mono if stereo
-                if audio.channels == 2:
-                    samples = samples.reshape((-1, 2)).mean(axis=1)
-                
-                # Convert to float32 numpy array
-                waveform = torch.tensor(samples).float() / 32768.0  # Assuming 16-bit audio
-                waveform = waveform.unsqueeze(0)  # Add channel dimension
-                
-                # Perform diarization
-                diarization = pipeline({"waveform": waveform, "sample_rate": audio.frame_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
+# Gradio Interface
+def gradio_interface(zip_file):
+    speaker1_zip, speaker2_zip = classify_and_group_speakers(zip_file)
+    return speaker1_zip, speaker2_zip
 
-iface = gr.Interface(
-    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."
-)
+gradio_inputs = gr.inputs.File(label="Upload ZIP of Audio Files", file_count="single")
+gradio_outputs = [gr.outputs.File(label="Speaker 1 ZIP"), gr.outputs.File(label="Speaker 2 ZIP")]
 
-iface.launch()
+gr.Interface(fn=gradio_interface, inputs=gradio_inputs, outputs=gradio_outputs, title="Speaker Diarization & Grouping").launch()