Update app.py

app.py

@@ -1,8 +1,9 @@
 import torch
 import os
+import gradio as gr
 from pyannote.audio import Pipeline
 from pydub import AudioSegment
-import gradio as gr
+from spaces import GPU
 
 # 获取 Hugging Face 认证令牌
 HF_TOKEN = os.environ.get("HUGGINGFACE_READ_TOKEN")
@@ -59,6 +60,7 @@ def combine_audio_with_time(target_audio, mixed_audio):
     return {"start_time": target_start_time, "end_time": target_end_time}
 
 # 使用 pyannote/speaker-diarization 对拼接后的音频进行说话人分离
+@GPU(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
 def diarize_audio(temp_file):
     if pipeline is None:
         return "错误: 模型未初始化"
@@ -97,7 +99,33 @@ def find_best_matching_speaker(target_start_time, target_end_time, diarization):
     
     return best_match, max_overlap
 
-# 获取最匹配的说话人并返回其时间段
+# 获取目标说话人的时间段（排除目标音频时间段）
+def get_speaker_segments(diarization, target_start_time, target_end_time, final_audio_length):
+    speaker_segments = {}
+    
+    # 遍历所有说话人时间段
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        start = turn.start
+        end = turn.end
+        
+        # 如果时间段与目标音频有重叠，需要截断
+        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))
+    
+    return speaker_segments
+
+# 处理音频文件并返回输出
 def process_audio(target_audio, mixed_audio):
     print(f"处理音频：目标音频: {target_audio}, 混合音频: {mixed_audio}")
     
@@ -125,10 +153,22 @@ def process_audio(target_audio, mixed_audio):
         )
         
         if best_match:
-            return {
-                'best_matching_speaker': best_match,
-                'overlap_duration': overlap_duration
-            }
+            # 获取目标说话人的时间段（排除和截断目标音频时间段）
+            speaker_segments = get_speaker_segments(
+                diarization_result, 
+                time_dict['start_time'], 
+                time_dict['end_time'], 
+                final_audio_length
+            )
+            
+            if best_match in speaker_segments:
+                return {
+                    'best_matching_speaker': best_match,
+                    'overlap_duration': overlap_duration,
+                    'segments': speaker_segments[best_match]
+                }
+            else:
+                return "没有找到匹配的说话人时间段。"
         else:
             return "未找到匹配的说话人。"
 
@@ -137,7 +177,7 @@ with gr.Blocks() as demo:
     gr.Markdown("""
     # 🗣️ 音频拼接与说话人分类 🗣️
     上传目标音频和混合音频，拼接并进行说话人分类。
-    结果包括最匹配的说话人以及重叠时长。
+    结果包括最匹配的说话人及其时间段，已排除和截断目标录音时间段。
     """)
     
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")
@@ -146,7 +186,7 @@ with gr.Blocks() as demo:
     process_button = gr.Button("处理音频")
     
     # 输出结果
-    diarization_output = gr.Textbox(label="最匹配的说话人及重叠时长")
+    diarization_output = gr.Textbox(label="说话人时间段")
 
     # 点击按钮时触发处理音频
     process_button.click(