Update app.py

app.py

@@ -1,9 +1,8 @@
 import torch
-import spaces
-import gradio as gr
 import os
 from pyannote.audio import Pipeline
 from pydub import AudioSegment
+import gradio as gr
 
 # 获取 Hugging Face 认证令牌
 HF_TOKEN = os.environ.get("HUGGINGFACE_READ_TOKEN")
@@ -60,7 +59,6 @@ def combine_audio_with_time(target_audio, mixed_audio):
     return {"start_time": target_start_time, "end_time": target_end_time}
 
 # 使用 pyannote/speaker-diarization 对拼接后的音频进行说话人分离
-@spaces.GPU(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
 def diarize_audio(temp_file):
     if pipeline is None:
         return "错误: 模型未初始化"
@@ -74,37 +72,34 @@ def diarize_audio(temp_file):
     except Exception as e:
         return f"处理音频时出错: {e}"
 
-# 获取指定说话人的时间段（排除目标音频时间段）
-def get_speaker_segments(diarization, speaker_name, target_start_time, target_end_time, final_audio_length):
-    speaker_segments = {}
+# 查找最匹配的说话人
+def find_best_matching_speaker(target_start_time, target_end_time, diarization):
+    best_match = None
+    max_overlap = 0
     
-    # 遍历所有说话人时间段
+    # 遍历所有说话人时间段，计算与目标音频的重叠部分
     for turn, _, speaker in diarization.itertracks(yield_label=True):
         start = turn.start
         end = turn.end
         
-        # 如果是目标说话人
-        if speaker == speaker_name:
-            # 如果时间段与目标音频有重叠，需要截断
-            if start < target_end_time and end > target_start_time:
-                # 记录被截断的时间段
-                if start < target_start_time:
-                    # 目标音频开始前的时间段
-                    speaker_segments.setdefault(speaker, []).append((start, min(target_start_time, end)))
-                
-                if end > target_end_time:
-                    # 目标音频结束后的时间段
-                    speaker_segments.setdefault(speaker, []).append((max(target_end_time, start), min(end, final_audio_length)))
-            else:
-                # 完全不与目标音频重叠的时间段
-                if end <= target_start_time or start >= target_end_time:
-                    speaker_segments.setdefault(speaker, []).append((start, end))
+        # 计算重叠部分的开始和结束时间
+        overlap_start = max(start, target_start_time)
+        overlap_end = min(end, target_end_time)
+        
+        # 如果有重叠部分，计算重叠的持续时间
+        if overlap_end > overlap_start:
+            overlap_duration = overlap_end - overlap_start
+            
+            # 如果当前重叠部分更大，则更新最匹配的说话人
+            if overlap_duration > max_overlap:
+                max_overlap = overlap_duration
+                best_match = speaker
     
-    return speaker_segments
+    return best_match, max_overlap
 
-# 处理音频文件并返回输出
-def process_audio(target_audio, mixed_audio, speaker_name):
-    print(f"处理音频：目标音频: {target_audio}, 混合音频: {mixed_audio}, 提取说话人: {speaker_name}")
+# 获取最匹配的说话人并返回其时间段
+def process_audio(target_audio, mixed_audio):
+    print(f"处理音频：目标音频: {target_audio}, 混合音频: {mixed_audio}")
     
     # 进行音频拼接并返回目标音频的起始和结束时间（作为字典）
     time_dict = combine_audio_with_time(target_audio, mixed_audio)
@@ -122,45 +117,41 @@ def process_audio(target_audio, mixed_audio, speaker_name):
         # 获取拼接后的音频长度
         final_audio_length = len(AudioSegment.from_wav("final_output.wav")) / 1000  # 秒为单位
         
-        # 获取目标说话人的时间段（排除目标音频时间段）
-        speaker_segments = get_speaker_segments(
-            diarization_result, 
-            speaker_name, 
+        # 查找最匹配的说话人
+        best_match, overlap_duration = find_best_matching_speaker(
             time_dict['start_time'], 
             time_dict['end_time'], 
-            final_audio_length
+            diarization_result
         )
         
-        if speaker_segments and speaker_name in speaker_segments:
-            # 返回目标说话人的时间段（已排除和截断目标音频时间段）
+        if best_match:
             return {
-                'segments': speaker_segments[speaker_name],
-                'total_duration': sum(end - start for start, end in speaker_segments[speaker_name])
+                'best_matching_speaker': best_match,
+                'overlap_duration': overlap_duration
             }
         else:
-            return f"没有找到 {speaker_name} 的时间段。"
+            return "未找到匹配的说话人。"
 
 # Gradio 接口
 with gr.Blocks() as demo:
-    gr.Markdown(""" 
+    gr.Markdown("""
     # 🗣️ 音频拼接与说话人分类 🗣️
     上传目标音频和混合音频，拼接并进行说话人分类。
-    ���果包括指定说话人的时间段，已排除和截断目标录音时间段。
+    结果包括最匹配的说话人以及重叠时长。
     """)
     
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")
     target_audio_input = gr.Audio(type="filepath", label="上传目标说话人音频")
-    speaker_name_input = gr.Textbox(label="请输入说话人名称（如 'SPEAKER_01'）", value="SPEAKER_00")
     
     process_button = gr.Button("处理音频")
     
     # 输出结果
-    diarization_output = gr.Textbox(label="说话人时间段")
+    diarization_output = gr.Textbox(label="最匹配的说话人及重叠时长")
 
     # 点击按钮时触发处理音频
     process_button.click(
         fn=process_audio,
-        inputs=[target_audio_input, mixed_audio_input, speaker_name_input],
+        inputs=[target_audio_input, mixed_audio_input],
         outputs=[diarization_output]
     )