Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+import spaces
 import gradio as gr
 import os
 from pyannote.audio import Pipeline
@@ -59,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 对拼接后的音频进行说话人分离
-@gr.Interface(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
+@spaces.GPU(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
 def diarize_audio(temp_file):
     if pipeline is None:
         return "错误: 模型未初始化"
@@ -73,8 +74,8 @@ def diarize_audio(temp_file):
     except Exception as e:
         return f"处理音频时出错: {e}"
 
-# 获取目标说话人的时间段并替换指定的SPEAKER_00
-def get_speaker_segments(diarization, target_start_time, target_end_time, final_audio_length):
+# 获取指定说话人的时间段（排除目标音频时间段）
+def get_speaker_segments(diarization, speaker_name, target_start_time, target_end_time, final_audio_length):
     speaker_segments = {}
     
     # 遍历所有说话人时间段
@@ -83,24 +84,27 @@ def get_speaker_segments(diarization, target_start_time, target_end_time, final_
         end = turn.end
         
         # 如果是目标说话人
-        if speaker == 'SPEAKER_00':
-            # 替换目标音频的时间段
+        if speaker == speaker_name:
+            # 如果时间段与目标音频有重叠，需要截断
             if start < target_end_time and end > target_start_time:
-                # 目标音频时间段被截断，重新计算其时间段
-                new_start = max(start, target_start_time)
-                new_end = min(end, target_end_time)
-                speaker_segments.setdefault(speaker, []).append((new_start, new_end))
-        
-        else:
-            # 完全不与目标音频重叠的时间段
-            if end <= target_start_time or start >= target_end_time:
-                speaker_segments.setdefault(speaker, []).append((start, end))
+                # 记录被截断的时间段
+                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}")
+def process_audio(target_audio, mixed_audio, speaker_name):
+    print(f"处理音频：目标音频: {target_audio}, 混合音频: {mixed_audio}, 提取说话人: {speaker_name}")
     
     # 进行音频拼接并返回目标音频的起始和结束时间（作为字典）
     time_dict = combine_audio_with_time(target_audio, mixed_audio)
@@ -118,33 +122,35 @@ 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, 
+            speaker_name, 
             time_dict['start_time'], 
             time_dict['end_time'], 
             final_audio_length
         )
         
-        if speaker_segments and 'SPEAKER_00' in speaker_segments:
+        if speaker_segments and speaker_name in speaker_segments:
             # 返回目标说话人的时间段（已排除和截断目标音频时间段）
             return {
-                'segments': speaker_segments['SPEAKER_00'],
-                'total_duration': sum(end - start for start, end in speaker_segments['SPEAKER_00'])
+                'segments': speaker_segments[speaker_name],
+                'total_duration': sum(end - start for start, end in speaker_segments[speaker_name])
             }
         else:
-            return "没有找到SPEAKER_00的时间段。"
+            return f"没有找到 {speaker_name} 的时间段。"
 
 # Gradio 接口
 with gr.Blocks() as demo:
-    gr.Markdown("""
+    gr.Markdown(""" 
     # 🗣️ 音频拼接与说话人分类 🗣️
     上传目标音频和混合音频，拼接并进行说话人分类。
-    结果包括目标说话人（SPEAKER_00）的时间段，已排除和截断目标录音时间段。
+    结果包括指定说话人的时间段，已排除和截断目标录音时间段。
     """)
     
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")
     target_audio_input = gr.Audio(type="filepath", label="上传目标说话人音频")
+    speaker_name_input = gr.Textbox(label="请输入说话人名称（如 'SPEAKER_01'）", value="SPEAKER_00")
     
     process_button = gr.Button("处理音频")
     
@@ -154,7 +160,7 @@ with gr.Blocks() as demo:
     # 点击按钮时触发处理音频
     process_button.click(
         fn=process_audio,
-        inputs=[target_audio_input, mixed_audio_input],
+        inputs=[target_audio_input, mixed_audio_input, speaker_name_input],
         outputs=[diarization_output]
     )