Update app.py

app.py

@@ -64,6 +64,8 @@ def diarize_audio(temp_file):
         diarization = pipeline(temp_file)
     except Exception as e:
         return f"处理音频时出错: {e}"
+
+    print(diarization)
     
     # 返回 diarization 类对象
     return diarization
@@ -84,58 +86,31 @@ def calculate_overlap(start1, end1, start2, end2):
     overlap_duration = max(0, overlap_end - overlap_start)
     return overlap_duration
 
-# 获取所有说话人时间段，排除目标音频时间段
+# 获取目标时间段和说话人时间段的重叠比例
 def get_matching_segments(target_time, diarization_output):
     target_start_time = target_time['start_time']
     target_end_time = target_time['end_time']
     
-    # 记录每个说话人与目标音频的重叠时间
-    speaker_overlap = {}
-
-    for speech_turn in diarization_output.itertracks(yield_label=True):
+    # 获取该说话人时间段的信息，排除目标录音时间段
+    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
         label = speech_turn[1]
         
-        # 计算目标音频时间段与该说话人时间段的重叠时间
+        # 计算目标音频时间段和说话人时间段的重叠时间
         overlap = calculate_overlap(target_start_time, target_end_time, start_seconds, end_seconds)
         
+        # 如果存在重叠，排除目标音频时间段
         if overlap > 0:
-            if label not in speaker_overlap:
-                speaker_overlap[label] = 0
-            speaker_overlap[label] += overlap
-
-    # 找到与目标音频时间段重叠最多的说话人
-    max_overlap_speaker = max(speaker_overlap, key=speaker_overlap.get, default=None)
-    
-    if max_overlap_speaker is None:
-        return "没有找到匹配的说话人"
-    
-    # 获取该说话人的所有时间段，排除目标音频的时间段
-    speaker_segments = []
-    for speech_turn in diarization_output.itertracks(yield_label=True):
-        start_seconds = speech_turn[0].start
-        end_seconds = speech_turn[0].end
-        label = speech_turn[1]
-        
-        if label == max_overlap_speaker:
-            # 计算目标音频时间段与该说话人时间段的重叠时间
-            overlap = calculate_overlap(target_start_time, target_end_time, start_seconds, end_seconds)
-            if overlap == 0:  # 如果没有重叠，则保留该时间段
-                speaker_segments.append((start_seconds, end_seconds))
-
-    # 转换时间段为更易读的格式（例如：00:00:03.895 --> 00:00:04.367）
-    formatted_segments = [
-        f"{format_time(segment[0])} --> {format_time(segment[1])}" for segment in speaker_segments
-    ]
-    
-    return formatted_segments
+            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))
 
-# 格式化时间（秒 -> hh:mm:ss.xxx）
-def format_time(seconds):
-    mins, secs = divmod(seconds, 60)
-    hrs, mins = divmod(mins, 60)
-    return f"{int(hrs):02}:{int(mins):02}:{secs:06.3f}"
+    return speaker_segments
 
 # 处理音频文件并返回输出
 def process_audio(target_audio, mixed_audio):
@@ -151,12 +126,12 @@ def process_audio(target_audio, mixed_audio):
     if isinstance(diarization_result, str) and diarization_result.startswith("错误"):
         return diarization_result, None  # 出错时返回错误信息
     else:
-        # 获取重叠最多的说话人的所有匹配时间段
+        # 获取该说话人的所有匹配时间段（排除目标音频时间段）
         matching_segments = get_matching_segments(time_dict, diarization_result)
         
         if matching_segments:
-            # 返回匹配的时间段
-            return "\n".join(matching_segments)
+            # 返回匹配的说话人标签和他们的时间段
+            return matching_segments
         else:
             return "没有找到匹配的说话人时间段。"
 
@@ -164,7 +139,7 @@ def process_audio(target_audio, mixed_audio):
 with gr.Blocks() as demo:
     gr.Markdown("""
     # 🗣️ 音频拼接与说话人分类 🗣️
-    上传目标音频和混合音频，拼接并进行说话人分类。结果包括匹配重叠最多的说话人的时间段。
+    上传目标音频和混合音频，拼接并进行说话人分类。结果包括所有匹配说话人的时间段（排除目标录音时间段）。
     """)
     
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")