Update app.py

app.py

@@ -1,9 +1,9 @@
 import torch
-import spaces
-import gradio as gr
 import os
+import gradio as gr
 from pyannote.audio import Pipeline
 from pydub import AudioSegment
+from spaces import GPU
 
 # 获取 Hugging Face 认证令牌
 HF_TOKEN = os.environ.get("HUGGINGFACE_READ_TOKEN")
@@ -60,7 +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 对拼接后的音频进行说话人分离
-@spaces.GPU(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
+@GPU(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
 def diarize_audio(temp_file):
     if pipeline is None:
         return "错误: 模型未初始化"
@@ -74,6 +74,31 @@ def diarize_audio(temp_file):
     except Exception as e:
         return f"处理音频时出错: {e}"
 
+# 查找最匹配的说话人
+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
+        
+        # 计算重叠部分的开始和结束时间
+        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 best_match, max_overlap
+
 # 获取目标说话人的时间段（排除目标音频时间段）
 def get_speaker_segments(diarization, target_start_time, target_end_time, final_audio_length):
     speaker_segments = {}
@@ -83,34 +108,23 @@ def get_speaker_segments(diarization, target_start_time, target_end_time, final_
         start = turn.start
         end = turn.end
         
-        # 如果是目标说话人
-        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:
-                # 完全不与目标音频重叠的时间段
-                if end <= target_start_time or start >= target_end_time:
-                    speaker_segments.setdefault(speaker, []).append((start, 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 clip_audio(audio_segment, segments):
-    clips = []
-    for start, end in segments:
-        start_ms = int(start * 1000)  # 毫秒
-        end_ms = int(end * 1000)  # 毫秒
-        clips.append(audio_segment[start_ms:end_ms])
-    return clips
-
 # 处理音频文件并返回输出
 def process_audio(target_audio, mixed_audio):
     print(f"处理音频：目标音频: {target_audio}, 混合音频: {mixed_audio}")
@@ -131,30 +145,44 @@ def process_audio(target_audio, mixed_audio):
         # 获取拼接后的音频长度
         final_audio_length = len(AudioSegment.from_wav("final_output.wav")) / 1000  # 秒为单位
         
-        # 获取目标说话人的时间段（排除和截断目标音频时间段）
-        speaker_segments = get_speaker_segments(
-            diarization_result, 
+        # 查找最匹配的说话人
+        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_00' in speaker_segments:
-            # 剪辑目标说话人的音频片段
-            final_audio_segment = AudioSegment.from_wav("final_output.wav")
-            clips = clip_audio(final_audio_segment, speaker_segments['SPEAKER_00'])
-            
-            # 将剪辑后的音频片段导出为多个文件
-            output_files = []
-            for i, clip in enumerate(clips):
-                clip_path = f"speaker_00_clip_{i + 1}.wav"
-                clip.export(clip_path, format="wav")
-                output_files.append(clip_path)
+        if best_match:
+            # 获取目标说话人的时间段（排除和截断目标音频时间段）
+            speaker_segments = get_speaker_segments(
+                diarization_result, 
+                time_dict['start_time'], 
+                time_dict['end_time'], 
+                final_audio_length
+            )
             
-            # 返回剪辑后的音频文件路径
-            return output_files
+            if best_match in speaker_segments:
+                # 拼接所有片段
+                final_output = AudioSegment.empty()
+                for segment in speaker_segments[best_match]:
+                    start_time_ms = int(segment[0] * 1000)  # 转为毫秒
+                    end_time_ms = int(segment[1] * 1000)
+                    segment_audio = AudioSegment.from_wav("final_output.wav")[start_time_ms:end_time_ms]
+                    final_output += segment_audio
+                
+                # 导出最终拼接音频
+                final_output.export("final_combined_output.wav", format="wav")
+                
+                return {
+                    'best_matching_speaker': best_match,
+                    'overlap_duration': overlap_duration,
+                    'segments': speaker_segments[best_match],
+                    'final_audio': "final_combined_output.wav"
+                }
+            else:
+                return "没有找到匹配的说话人时间段。"
         else:
-            return "没有找到SPEAKER_00的时间段。"
+            return "未找到匹配的说话人。"
 
 # Gradio 接口
 with gr.Blocks() as demo: