Update app.py

app.py

@@ -20,6 +20,11 @@ except Exception as e:
     print(f"Error initializing pipeline: {e}")
     pipeline = None
 
+# 时间戳转换为秒
+def timestamp_to_seconds(timestamp):
+    h, m, s = map(float, timestamp.split(':'))
+    return 3600 * h + 60 * m + s
+
 # 音频拼接函数：拼接目标音频和混合音频，返回目标音频的起始时间和结束时间作为字典
 def combine_audio_with_time(target_audio, mixed_audio):
     if pipeline is None:
@@ -62,57 +67,43 @@ def diarize_audio(temp_file):
     
     try:
         diarization = pipeline(temp_file)
+        print("说话人分离结果:")
+        for turn, _, speaker in diarization.itertracks(yield_label=True):
+            print(f"[{turn.start:.3f} --> {turn.end:.3f}] {speaker}")
+        return diarization
     except Exception as e:
         return f"处理音频时出错: {e}"
 
-    print(diarization)
-    
-    # 返回 diarization 类对象
-    return diarization
-
-# 获取目标录音所在时间范围最大的说话人及其时间段
-def get_most_matched_speaker_segments(diarization_output, target_start_time, target_end_time, final_audio_length):
-    # 用于存储说话人与目标音频重叠时间的字典
-    speaker_overlaps = {}
+# 获取目标说话人的时间段（排除目标音频时间段）
+def get_speaker_segments(diarization, target_start_time, target_end_time, final_audio_length):
+    speaker_segments = {}
     
-    # 使用 itertracks 获取每个说话人的信息
-    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]
+    # 遍历所有说话人时间段
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        start = turn.start
+        end = turn.end
         
-        # 计算目标音频与当前说话人时间段的重叠时间
-        overlap_start = max(start_seconds, target_start_time)
-        overlap_end = min(end_seconds, target_end_time)
-        overlap_duration = max(0, overlap_end - overlap_start)
-        
-        # 如果有重叠，记录重叠时间
-        if overlap_duration > 0:
-            if label not in speaker_overlaps:
-                speaker_overlaps[label] = {
-                    'total_overlap': overlap_duration,
-                    'segments': []
-                }
+        # 如果是目标说话人
+        if speaker == 'SPEAKER_00':
+            # 如果时间段与目标音频有重叠，需要截断
+            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:
-                speaker_overlaps[label]['total_overlap'] += overlap_duration
-            
-            # 记录该说话人的原始时间段（排除目标音频时间段）
-            if start_seconds < target_start_time:
-                speaker_overlaps[label]['segments'].append((start_seconds, min(end_seconds, target_start_time)))
-            
-            if end_seconds > target_end_time:
-                speaker_overlaps[label]['segments'].append((max(start_seconds, target_end_time), end_seconds))
-    
-    # 找到重叠时间最长的说话人
-    if speaker_overlaps:
-        most_matched_speaker = max(speaker_overlaps, key=lambda k: speaker_overlaps[k]['total_overlap'])
-        return {most_matched_speaker: speaker_overlaps[most_matched_speaker]['segments']}
+                # 完全不与目标音频重叠的时间段
+                if end <= target_start_time or start >= target_end_time:
+                    speaker_segments.setdefault(speaker, []).append((start, end))
     
-    return {}
+    return speaker_segments
 
 # 处理音频文件并返回输出
 def process_audio(target_audio, mixed_audio):
-    # 打印文件路径，确保传入的文件有效
     print(f"处理音频：目标音频: {target_audio}, 混合音频: {mixed_audio}")
     
     # 进行音频拼接并返回目标音频的起始和结束时间（作为字典）
@@ -131,25 +122,29 @@ def process_audio(target_audio, mixed_audio):
         # 获取拼接后的音频长度
         final_audio_length = len(AudioSegment.from_wav("final_output.wav")) / 1000  # 秒为单位
         
-        # 获取目标录音所在时间范围最大的说话人时间段
-        most_matched_speaker_segments = get_most_matched_speaker_segments(
+        # 获取目标说话人的时间段（排除目标音频时间段）
+        speaker_segments = get_speaker_segments(
             diarization_result, 
             time_dict['start_time'], 
             time_dict['end_time'], 
             final_audio_length
         )
         
-        if most_matched_speaker_segments:
-            # 返回目标录音所在时间范围最大的说话人的时间段（排除目标音频时间段）
-            return most_matched_speaker_segments
+        if speaker_segments and 'SPEAKER_00' in speaker_segments:
+            # 返回目标说话人的时间段（已排除和截断目标音频时间段）
+            return {
+                'segments': speaker_segments['SPEAKER_00'],
+                'total_duration': sum(end - start for start, end in speaker_segments['SPEAKER_00'])
+            }
         else:
-            return "没有找到与目标录音重叠的说话人时间段。"
+            return "没有找到SPEAKER_00的时间段。"
 
 # Gradio 接口
 with gr.Blocks() as demo:
     gr.Markdown("""
     # 🗣️ 音频拼接与说话人分类 🗣️
-    上传目标音频和混合音频，拼接并进行说话人分类。结果包括与目标录音重叠时间最长的说话人的时间段（排除目标录音时间段）。
+    上传目标音频和混合音频，拼接并进行说话人分类。
+    结果包括目标说话人（SPEAKER_00）的时间段，已排除和截断目标录音时间段。
     """)
     
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")