Update app.py

app.py

@@ -1,90 +1,87 @@
-import os
+import gradio as gr
 import zipfile
+import os
 import shutil
-import torch
-import torchaudio
 from pyannote.audio import Pipeline
-import gradio as gr
+import torch
 
-# Load the pre-trained model using your Hugging Face access token
-HUGGINGFACE_ACCESS_TOKEN = os.getenv("HF_TOKEN")
+# Set up the directory for processing
+TEMP_DIR = "temp_audio"
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+# Initialize the pyannote.audio pipeline
 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)
+    use_auth_token=os.getenv("HF_TOKEN")
+)
+
+# Move pipeline to GPU if available
+if torch.cuda.is_available():
+    pipeline.to(torch.device("cuda"))
+
+def process_audio_zip(file_info):
+    # Unzip the uploaded file
+    with zipfile.ZipFile(file_info, 'r') as zip_ref:
+        zip_ref.extractall(TEMP_DIR)
+    
+    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: 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)
-        
-        # Convert MP3 to waveform
-        waveform, sample_rate = torchaudio.load(audio_fp)
-
-        # Ensure the audio is mono and at 16 kHz
-        if waveform.shape[0] > 1:
-            waveform = waveform.mean(dim=0).unsqueeze(0)
-        if sample_rate != 16000:
-            waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
-            sample_rate = 16000
-
-        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)
-
-    # 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)
-
-    # Step 5: Clean up temporary directories
-    shutil.rmtree(extract_dir)
-    shutil.rmtree(speaker1_dir)
-    shutil.rmtree(speaker2_dir)
-
+    # Process each audio file in the temporary directory
+    for filename in os.listdir(TEMP_DIR):
+        if filename.endswith(".wav"):
+            file_path = os.path.join(TEMP_DIR, filename)
+            
+            # Run the diarization pipeline
+            diarization = pipeline(file_path)
+            
+            # Determine if the audio is mostly from speaker1 or speaker2
+            total_duration = {1: 0.0, 2: 0.0}
+            for turn, _, speaker in diarization.itertracks(yield_label=True):
+                total_duration[speaker] += turn.duration
+            
+            # Move file to the corresponding speaker directory
+            dominant_speaker = 1 if total_duration[1] >= total_duration[2] else 2
+            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))
+
+    # Zip the results
+    speaker1_zip = "speaker1.zip"
+    speaker2_zip = "speaker2.zip"
+    
+    def zipdir(path, ziph):
+        # Zip the directories
+        for root, dirs, files in os.walk(path):
+            for file in files:
+                ziph.write(os.path.join(root, file),
+                           os.path.relpath(os.path.join(root, file),
+                           os.path.join(path, '..')))
+
+    with zipfile.ZipFile(speaker1_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        zipdir(speaker1_dir, zipf)
+    
+    with zipfile.ZipFile(speaker2_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        zipdir(speaker2_dir, zipf)
+    
+    # Clean up the temporary directory
+    shutil.rmtree(TEMP_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
-
-gradio_inputs = gr.File(label="Upload ZIP of Audio Files", file_count="single")
-gradio_outputs = [gr.File(label="Speaker 1 ZIP"), gr.File(label="Speaker 2 ZIP")]
-
-gr.Interface(fn=gradio_interface, inputs=gradio_inputs, outputs=gradio_outputs, title="Speaker Diarization & Grouping").launch()
+# Gradio interface
+iface = gr.Interface(
+    fn=process_audio_zip,
+    inputs=gr.inputs.File(type="file"),
+    outputs=[
+        gr.File(label="Speaker 1 Audio"),
+        gr.File(label="Speaker 2 Audio")
+    ],
+    title="Speaker Diarization",
+    description="Upload a ZIP file containing audio files, and this will return two ZIP files containing diarized audio for each speaker."
+)
+
+iface.launch()