Update app.py

app.py

@@ -64,7 +64,8 @@ def diarize_audio(temp_file):
         diarization = pipeline(temp_file)
     except Exception as e:
         return f"处理音频时出错: {e}"
-    
+
+    print(diarization)
     # 返回 diarization 类对象
     return diarization
 
@@ -85,12 +86,12 @@ def calculate_overlap(start1, end1, start2, end2):
     return overlap_duration
 
 # 获取目标时间段和说话人时间段的重叠比例
-def get_best_match(target_time, diarization_output):
+def get_matching_segments(target_time, diarization_output):
     target_start_time = target_time['start_time']
     target_end_time = target_time['end_time']
     
-    # 通过 diarization_output 获取说话人信息
-    speaker_segments = []
+    # 获取该说话人时间段的信息，排除目标录音时间段
+    speaker_segments = {}
     for speech_turn in diarization_output.itertracks(yield_label=True):  # 使用 itertracks 获取每个说话人的信息
         start_seconds = speech_turn[0].start
         end_seconds = speech_turn[0].end
@@ -98,15 +99,17 @@ def get_best_match(target_time, diarization_output):
         
         # 计算目标音频时间段和说话人时间段的重叠时间
         overlap = calculate_overlap(target_start_time, target_end_time, start_seconds, end_seconds)
-        overlap_ratio = overlap / (end_seconds - start_seconds)
         
-        # 记录说话人标签和重叠比例
-        speaker_segments.append((label, overlap_ratio, start_seconds, end_seconds))
+        # 如果存在重叠，排除目标音频时间段
+        if overlap > 0:
+            if label not in speaker_segments:
+                speaker_segments[label] = []
+            
+            # 如果时间段与目标音频有重叠，跳过该时间段
+            if start_seconds >= target_end_time or end_seconds <= target_start_time:
+                speaker_segments[label].append((start_seconds, end_seconds))
 
-    # 按照重叠比例排序，返回重叠比例最大的一段
-    best_match = max(speaker_segments, key=lambda x: x[1], default=None)
-    
-    return best_match
+    return speaker_segments
 
 # 处理音频文件并返回输出
 def process_audio(target_audio, mixed_audio):
@@ -122,18 +125,20 @@ def process_audio(target_audio, mixed_audio):
     if isinstance(diarization_result, str) and diarization_result.startswith("错误"):
         return diarization_result, None  # 出错时返回错误信息
     else:
-        # 获取最佳匹配的说话人标签和时间段
-        best_match = get_best_match(time_dict, diarization_result)
+        # 获取该说话人的所有匹配时间段（排除目标音频时间段）
+        matching_segments = get_matching_segments(time_dict, diarization_result)
         
-        if best_match:
-            # 返回最佳匹配说话人的标签和时间段
-            return best_match[0], best_match[2], best_match[3]
+        if matching_segments:
+            # 返回匹配的说话人标签和他们的时间段
+            return matching_segments
+        else:
+            return "没有找到匹配的说话人时间段。"
 
 # Gradio 接口
 with gr.Blocks() as demo:
     gr.Markdown("""
     # 🗣️ 音频拼接与说话人分类 🗣️
-    上传目标音频和混合音频，拼接并进行说话人分类。结果包括最佳匹配说话人的时间段。
+    上传目标音频和混合音频，拼接并进行说话人分类。结果包括所有匹配说话人的时间段（排除目标录音时间段）。
     """)
     
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")
@@ -142,14 +147,13 @@ with gr.Blocks() as demo:
     process_button = gr.Button("处理音频")
     
     # 输出结果
-    diarization_output = gr.Textbox(label="最佳匹配说话人")
-    time_range_output = gr.Textbox(label="最佳匹配时间段")
+    diarization_output = gr.Textbox(label="匹配的说话人时间段")
 
     # 点击按钮时触发处理音频
     process_button.click(
         fn=process_audio,
         inputs=[target_audio_input, mixed_audio_input],
-        outputs=[diarization_output, time_range_output]
+        outputs=[diarization_output]
     )
 
 demo.launch(share=True)