音频中双人声音分离

import soundfile as sf
import librosa
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import numpy as np

# 1. 初始化语音分离 pipeline
print("正在加载模型...")
separation_pipeline = pipeline(
    Tasks.speech_separation,
    model='damo/speech_mossformer2_separation_temporal_8k'
)

# 2. 准备音频文件
input_audio_path = '1.6.WAV' # 这里改成你的文件名
temp_8k_file = 'temp_input_8k.wav'  # 临时文件

print(f"正在读取并重采样音频: {input_audio_path} ...")

# === 关键步骤：使用 librosa 读取并强制重采样为 8000Hz ===
# sr=8000 会自动把 44.1k/48k 的音频转为 8k
audio_data, sr = librosa.load(input_audio_path, sr=8000)

# 保存一个临时的 8k 文件给模型读取（这样最稳妥，防止格式报错）
sf.write(temp_8k_file, audio_data, 8000)

# 3. 执行分离
print("正在进行盲源分离 (这可能需要一点时间)...")
result = separation_pipeline(temp_8k_file)

# # result['output_pcm_list'] 包含分离后的两个音频数据
# # 这里的关键是：先取出来，然后用 .squeeze() 去掉多余的维度
# audio_track_1 = result['output_pcm_list'][0].squeeze()
# audio_track_2 = result['output_pcm_list'][1].squeeze()
# === 修复重点：将 bytes 转为 numpy 数组 ===
raw_bytes_1 = result['output_pcm_list'][0]
raw_bytes_2 = result['output_pcm_list'][1]

print(f"获取到数据，数据类型: {type(raw_bytes_1)}")
print(f"数据长度: {len(raw_bytes_1)} 字节")

# 使用 np.frombuffer 将字节流转为数组
# ModelScope 的 PCM 输出通常是 16-bit 整数 (int16)
audio_track_1 = np.frombuffer(raw_bytes_1, dtype=np.int16)
audio_track_2 = np.frombuffer(raw_bytes_2, dtype=np.int16)

# === 调试信息：打印一下形状 ===
print(f"音轨1 形状: {audio_track_1.shape}") # 正常应该是 (N,) 比如 (80000,)
print(f"音轨2 形状: {audio_track_2.shape}")

# 如果形状是 (Channels, Time) 例如 (1, 80000)，squeeze 会把它变成 (80000,)
# 如果形状是 (Time, Channels) 例如 (80000, 1)，squeeze 也会把它变成 (80000,)

# 4. 保存结果
print("正在保存分离后的音轨...")

# 这里的 .T 是为了防止万一 squeeze 没生效，数据是 (1, N) 的话，转置成 (N, 1) 比较保险
# 但通常 squeeze() 就足够了。我们加一个保险判断：
# if audio_track_1.ndim == 2 and audio_track_1.shape[0] < audio_track_1.shape[1]:
#     audio_track_1 = audio_track_1.T
# if audio_track_2.ndim == 2 and audio_track_2.shape[0] < audio_track_2.shape[1]:
#     audio_track_2 = audio_track_2.T

sf.write('output_speaker_A.wav', audio_track_1, 8000)
sf.write('output_speaker_B.wav', audio_track_2, 8000)

print("成功！已生成 output_speaker_A.wav 和 output_speaker_B.wav")

# ---------------------------------------------------------
# 5. (可选) 自动区分男女
# ---------------------------------------------------------
# def detect_gender(y, sr=8000):
#     # 提取音高
#     f0, voiced_flag, voiced_probs = librosa.pyin(y, fmin=50, fmax=300, sr=sr)
#     valid_f0 = f0[~np.isnan(f0)]
#     if len(valid_f0) == 0:
#         return 0
#     return np.mean(valid_f0)

def detect_gender(y, sr=8000):
    print(f"正在分析音频片段，数据类型: {y.dtype}")
    
    # === 关键修复：如果数据是整数 (int16)，强制转为浮点数 (float32) ===
    # librosa 要求输入必须是 float，且范围在 -1 到 1 之间
    if not np.issubdtype(y.dtype, np.floating):
        y = y.astype(np.float32) / 32768.0
    
    # 提取音高 (fmin=50, fmax=300 涵盖了男女生的正常语调范围)
    f0, voiced_flag, voiced_probs = librosa.pyin(y, fmin=50, fmax=300, sr=sr)
    
    # 去掉无效值 (nan)
    valid_f0 = f0[~np.isnan(f0)]
    
    if len(valid_f0) == 0:
        return 0
    return np.mean(valid_f0)

pitch_a = detect_gender(audio_track_1)
pitch_b = detect_gender(audio_track_2)

print(f"音轨A 平均频率: {pitch_a:.1f} Hz")
print(f"音轨B 平均频率: {pitch_b:.1f} Hz")

if pitch_a > pitch_b:
    # print("推测: A是女性, B是男性")
else:
    # print("推测: A是男性, B是女性")