# -*- coding: utf-8 -*- """ 说话人分离模块 基于 CAM++ 的说话人分离，用于多说话人音频分割 """ from loguru import logger import numpy as np import librosa import soundfile as sf import tempfile import os import threading from typing import Any, List, Mapping, Optional, Sequence, cast from dataclasses import dataclass import torch from ..core.config import settings from ..core.exceptions import DefaultServerErrorException # 全局 CAM++ pipeline 缓存（单例） _global_diarization_pipeline: Any | None = None _diarization_pipeline_lock = threading.Lock() _diarization_inference_semaphore = threading.BoundedSemaphore(1) @dataclass class SpeakerSegment: """说话人分段信息""" start_ms: int end_ms: int speaker_id: str audio_data: Optional[np.ndarray] = None temp_file: Optional[str] = None @property def start_sec(self) -> float: return self.start_ms / 1000.0 @property def end_sec(self) -> float: return self.end_ms / 1000.0 @property def duration_ms(self) -> int: return self.end_ms - self.start_ms @property def duration_sec(self) -> float: return self.duration_ms / 1000.0 def _resolve_modelscope_device() -> str: """根据配置和硬件自动选择 modelscope pipeline 设备 """ from ..core.device import detect_device return detect_device(settings.DEVICE) def _move_pipeline_model_to_device(pipeline_instance: Any, modelscope_device: str) -> None: """将 pipeline 的底层模型迁移到目标设备。""" if hasattr(pipeline_instance, "device_name"): pipeline_instance.device_name = modelscope_device model = getattr(pipeline_instance, "model", None) if model is not None and hasattr(model, "to"): pipeline_instance.model = model.to(modelscope_device) def _create_modelscope_pipeline( *, task: Any, model: str, modelscope_device: str, model_revision: Optional[str] = None, ) -> Any: """创建 modelscope pipeline，并在需要时把底层模型迁移到目标设备。""" from modelscope.pipelines import pipeline pipeline_kwargs: dict[str, Any] = { "task": task, "model": model, "device": modelscope_device, } if model_revision is not None: pipeline_kwargs["model_revision"] = model_revision pipeline_instance = pipeline(**pipeline_kwargs) _move_pipeline_model_to_device(pipeline_instance, modelscope_device) return pipeline_instance def _enable_batched_sv( pipeline_instance: Any, modelscope_device: str, max_batch_size: int = 32, ) -> Any: """ 对说话人分离 pipeline 启用 batched SV 推理。 原始 pipeline 的 forward 方法逐个 segment 调用 sv_pipeline 提取 embedding， 这里改为将所有 segment 拼成一个 batch 一次性推理，大幅减少 GPU 调用次数。 同时将子 pipeline（sv / vad / change_locator）绑定到指定 device。 Args: pipeline_instance: CAM++ diarization pipeline 实例 modelscope_device: 设备名称 max_batch_size: 最大批处理大小，防止 OOM """ if getattr(pipeline_instance, "_batched_sv_enabled", False): return pipeline_instance from modelscope.utils.constant import Tasks config = getattr(pipeline_instance, "config", None) if not isinstance(config, Mapping): logger.warning("CAM++ pipeline 缺少可读取的 config，跳过 batched SV 优化") return pipeline_instance sv_model = config.get("speaker_model") vad_model = config.get("vad_model") change_locator = config.get("change_locator") if isinstance(sv_model, str) and sv_model: pipeline_instance.sv_pipeline = _create_modelscope_pipeline( task=Tasks.speaker_verification, model=sv_model, modelscope_device=modelscope_device, ) if isinstance(vad_model, str) and vad_model: pipeline_instance.vad_pipeline = _create_modelscope_pipeline( task=Tasks.voice_activity_detection, model=vad_model, modelscope_device=modelscope_device, model_revision="v2.0.2", ) if isinstance(change_locator, str) and change_locator: pipeline_instance.change_locator_pipeline = _create_modelscope_pipeline( task=Tasks.speaker_diarization, model=change_locator, modelscope_device=modelscope_device, ) def batched_forward(self: Any, segments: Sequence[Sequence[Any]]) -> np.ndarray: """批量提取说话人 embedding，替代逐段串行推理""" sv_model_instance = getattr(getattr(self, "sv_pipeline", None), "model", None) emb_size = int(getattr(sv_model_instance, "emb_size", 192)) if not segments: return np.empty((0, emb_size), dtype=np.float32) if sv_model_instance is None: raise RuntimeError("CAM++ sv_pipeline.model 未初始化") all_embeddings: list[np.ndarray] = [] total_segments = len(segments) start_idx = 0 while start_idx < total_segments: end_idx = min(start_idx + max_batch_size, total_segments) batch_segments = segments[start_idx:end_idx] batch_items: list[np.ndarray] = [] for segment in batch_segments: if len(segment) < 3: continue batch_items.append(np.asarray(segment[2], dtype=np.float32)) if not batch_items: start_idx = end_idx continue batch = np.stack(batch_items, axis=0) with torch.no_grad(): embeddings = sv_model_instance( cast(Any, torch).as_tensor(batch).to(modelscope_device) ) if isinstance(embeddings, torch.Tensor): all_embeddings.append(embeddings.detach().cpu().numpy()) else: all_embeddings.append(np.asarray(embeddings, dtype=np.float32)) start_idx = end_idx if not all_embeddings: return np.empty((0, emb_size), dtype=np.float32) return ( np.concatenate(all_embeddings, axis=0) if len(all_embeddings) > 1 else all_embeddings[0] ) import types pipeline_instance.forward = types.MethodType(batched_forward, pipeline_instance) pipeline_instance._batched_sv_enabled = True logger.info( "CAM++ 说话人分离启用 batched SV: device={}, sv_device={}, vad_device={}", modelscope_device, getattr(getattr(pipeline_instance, "sv_pipeline", None), "device_name", "unknown"), getattr(getattr(pipeline_instance, "vad_pipeline", None), "device_name", "unknown"), ) return pipeline_instance def get_global_diarization_pipeline() -> Any: """获取全局说话人分离 pipeline（懒加载单例）""" global _global_diarization_pipeline with _diarization_pipeline_lock: if _global_diarization_pipeline is None: try: from modelscope.utils.constant import Tasks from ..infrastructure.model_utils import resolve_model_path model_id = 'iic/speech_campplus_speaker-diarization_common' model_path = resolve_model_path(model_id) modelscope_device = _resolve_modelscope_device() logger.info( "正在加载 CAM++ 说话人分离模型: {}, device={}", model_path, modelscope_device, ) _global_diarization_pipeline = _create_modelscope_pipeline( task=Tasks.speaker_diarization, model=model_path, modelscope_device=modelscope_device, ) _global_diarization_pipeline = _enable_batched_sv( _global_diarization_pipeline, modelscope_device ) logger.info("CAM++ 模型加载成功（已启用 batched SV）") except Exception as e: logger.error(f"CAM++ 模型加载失败: {e}") raise DefaultServerErrorException(f"说话人分离模型加载失败: {str(e)}") return _global_diarization_pipeline class SpeakerDiarizer: """基于 CAM++ 的说话人分离器""" DEFAULT_MIN_SEGMENT_SEC = 1.0 DEFAULT_SAMPLE_RATE = 16000 LOW_ENERGY_SEARCH_WINDOW_MS = 10000 LOW_ENERGY_CONTEXT_MS = 160 LOW_ENERGY_STEP_MS = 20 def __init__( self, min_segment_sec: float = DEFAULT_MIN_SEGMENT_SEC, ): self.min_segment_sec = min_segment_sec self.min_segment_ms = int(min_segment_sec * 1000) def diarize( self, audio_path: str ) -> List[SpeakerSegment]: """执行说话人分离 Args: audio_path: 音频文件路径 Returns: 原始分段列表（未合并） """ audio_duration_ms: Optional[int] = None try: try: audio_duration_ms = int(librosa.get_duration(path=audio_path) * 1000) except Exception: audio_duration_ms = None # CAM++ 对极短片段收益很低，且容易直接报 "too short"。 # 这里提前降级成单说话人，避免无意义 warning 刷屏。 if audio_duration_ms is not None and audio_duration_ms < self.min_segment_ms: logger.debug( "音频时长过短，跳过 CAM++ 说话人分离: duration_ms=%s < min_segment_ms=%s", audio_duration_ms, self.min_segment_ms, ) return [ SpeakerSegment( start_ms=0, end_ms=max(audio_duration_ms, 1), speaker_id="说话人1", ) ] pipeline = get_global_diarization_pipeline() logger.info(f"开始说话人分离: {audio_path}") with _diarization_inference_semaphore: result = pipeline(audio_path) # 解析结果: {'text': [[start, end, speaker_id], ...]} # pipeline 返回类型不确定，需要安全地获取 'text' 字段 if isinstance(result, dict): raw_output = result.get('text', []) else: raw_output = getattr(result, 'text', []) or [] segments = [] for seg in raw_output: if isinstance(seg, list) and len(seg) == 3: try: start_ms = int(float(seg[0]) * 1000) end_ms = int(float(seg[1]) * 1000) speaker_id = f"说话人{int(seg[2]) + 1}" segments.append(SpeakerSegment( start_ms=start_ms, end_ms=end_ms, speaker_id=speaker_id, )) except (ValueError, TypeError) as e: logger.warning(f"跳过格式错误的片段: {seg}, 错误: {e}") logger.info(f"说话人分离完成，原始片段数: {len(segments)}") # 诊断日志：打印前20个原始片段 for i, seg in enumerate(segments[:20]): logger.debug( f"[CAM++原始] #{i}: {seg.start_sec:.2f}-{seg.end_sec:.2f}s " f"({seg.duration_sec:.2f}s) {seg.speaker_id}" ) return segments except Exception as e: error_msg = str(e).lower() # 音频太短时，返回默认的单说话人片段 if "too short" in error_msg: logger.debug("CAM++ 跳过过短音频，回退单说话人片段: %s", e) if audio_duration_ms is None: try: audio_duration_ms = int(librosa.get_duration(path=audio_path) * 1000) except Exception: audio_duration_ms = 5000 return [ SpeakerSegment( start_ms=0, end_ms=audio_duration_ms, speaker_id="说话人1", ) ] # 其他异常正常抛出 logger.error(f"说话人分离失败: {e}") raise DefaultServerErrorException(f"说话人分离失败: {str(e)}") def merge_consecutive_segments( self, segments: List[SpeakerSegment] ) -> List[SpeakerSegment]: """合并同一说话人的连续片段""" if not segments: return [] # 按开始时间排序 sorted_segments = sorted(segments, key=lambda x: x.start_ms) merged = [] current = SpeakerSegment( start_ms=sorted_segments[0].start_ms, end_ms=sorted_segments[0].end_ms, speaker_id=sorted_segments[0].speaker_id, ) for seg in sorted_segments[1:]: if seg.speaker_id == current.speaker_id: # 同一说话人，扩展结束时间 current.end_ms = max(current.end_ms, seg.end_ms) else: # 不同说话人，保存当前段，开始新段 logger.debug( f"[合并中断] 说话人切换: {current.speaker_id} → {seg.speaker_id} " f"在 {seg.start_sec:.2f}s，保存片段 {current.start_sec:.2f}-{current.end_sec:.2f}s" ) merged.append(current) current = SpeakerSegment( start_ms=seg.start_ms, end_ms=seg.end_ms, speaker_id=seg.speaker_id, ) # 保存最后一段 merged.append(current) logger.info(f"合并同一说话人连续片段: {len(segments)} → {len(merged)}") # 诊断日志：打印合并后的前20个片段 for i, seg in enumerate(merged[:20]): logger.debug( f"[合并后] #{i}: {seg.start_sec:.2f}-{seg.end_sec:.2f}s " f"({seg.duration_sec:.2f}s) {seg.speaker_id}" ) return merged def merge_short_segments( self, segments: List[SpeakerSegment] ) -> List[SpeakerSegment]: """智能合并短片段 策略： 1. 第一层：<10s的片段向后合并（避免孤立短片段） 2. 第二层：60s累积合并（合并连续片段） """ if not segments: return [] max_segment_sec = settings.MAX_SEGMENT_SEC # 按开始时间排序 sorted_segments = sorted(segments, key=lambda x: x.start_ms) # 第一层：<10s累积向后合并（循环计算直到>=10s或超过60s） merged = [] i = 0 while i < len(sorted_segments): seg = sorted_segments[i] # 如果>=10s，直接添加 if seg.duration_sec >= 10.0: merged.append(seg) i += 1 continue # <10s，开始累积合并 current_start_ms = seg.start_ms current_end_ms = seg.end_ms current_duration_sec = seg.duration_sec j = i + 1 # 累积合并，只要<10s且同说话人且不超过60s while j < len(sorted_segments) and current_duration_sec < 10.0: next_seg = sorted_segments[j] if next_seg.speaker_id != seg.speaker_id: break new_duration = (next_seg.end_ms - current_start_ms) / 1000.0 if new_duration > max_segment_sec: break current_end_ms = next_seg.end_ms current_duration_sec = new_duration j += 1 # 创建合并后的片段 merged_seg = SpeakerSegment( start_ms=current_start_ms, end_ms=current_end_ms, speaker_id=seg.speaker_id, ) merged.append(merged_seg) if j > i + 1: logger.debug( f"[第一层] {seg.speaker_id}: " f"累积合并了 {j - i} 个片段，结果 {merged_seg.duration_sec:.1f}s" ) i = j # 第二层：60s累积合并 final_merged = [] i = 0 while i < len(merged): seg = merged[i] current_start_ms = seg.start_ms current_end_ms = seg.end_ms j = i + 1 # 累积合并，只要 <= 60s 且同说话人 while j < len(merged): next_seg = merged[j] if next_seg.speaker_id != seg.speaker_id: break new_duration = (next_seg.end_ms - current_start_ms) / 1000.0 if new_duration > max_segment_sec: break current_end_ms = next_seg.end_ms j += 1 merged_seg = SpeakerSegment( start_ms=current_start_ms, end_ms=current_end_ms, speaker_id=seg.speaker_id, ) final_merged.append(merged_seg) if j > i + 1: logger.debug( f"[第二层] {seg.speaker_id}: " f"合并了 {j - i} 个片段" ) i = j return final_merged def _find_low_energy_boundary_ms( self, audio_data: np.ndarray, sample_rate: int, lower_ms: int, upper_ms: int, ) -> int: lower_ms = max(0, lower_ms) upper_ms = max(lower_ms, upper_ms) if sample_rate <= 0 or audio_data.size == 0: return upper_ms context_samples = max( 1, int(sample_rate * self.LOW_ENERGY_CONTEXT_MS / 1000) ) candidate_points = list( range(lower_ms, upper_ms + 1, self.LOW_ENERGY_STEP_MS) ) if not candidate_points or candidate_points[-1] != upper_ms: candidate_points.append(upper_ms) best_ms = upper_ms best_energy = float("inf") audio_length = int(audio_data.shape[0]) for candidate_ms in candidate_points: center_sample = int(candidate_ms * sample_rate / 1000) start_sample = max(0, center_sample - context_samples // 2) end_sample = min(audio_length, center_sample + context_samples // 2) if start_sample >= end_sample: continue window = audio_data[start_sample:end_sample] energy = float(np.mean(np.square(window))) if energy <= best_energy: best_energy = energy best_ms = candidate_ms return best_ms def split_long_segments( self, segments: List[SpeakerSegment], audio_data: np.ndarray, sample_rate: int, ) -> List[SpeakerSegment]: max_segment_ms = int(settings.MAX_SEGMENT_SEC * 1000) if max_segment_ms <= 0: return segments split_segments: List[SpeakerSegment] = [] for seg in segments: if seg.duration_ms <= max_segment_ms: split_segments.append(seg) continue current_start_ms = seg.start_ms while seg.end_ms - current_start_ms > max_segment_ms: hard_boundary_ms = current_start_ms + max_segment_ms lower_boundary_ms = max( current_start_ms + self.min_segment_ms, hard_boundary_ms - self.LOW_ENERGY_SEARCH_WINDOW_MS, ) boundary_ms = self._find_low_energy_boundary_ms( audio_data=audio_data, sample_rate=sample_rate, lower_ms=lower_boundary_ms, upper_ms=hard_boundary_ms, ) if boundary_ms <= current_start_ms: boundary_ms = hard_boundary_ms split_segments.append( SpeakerSegment( start_ms=current_start_ms, end_ms=boundary_ms, speaker_id=seg.speaker_id, ) ) current_start_ms = boundary_ms remaining_ms = seg.end_ms - current_start_ms if remaining_ms >= self.min_segment_ms: split_segments.append( SpeakerSegment( start_ms=current_start_ms, end_ms=seg.end_ms, speaker_id=seg.speaker_id, ) ) elif split_segments: split_segments[-1].end_ms = seg.end_ms if len(split_segments) != len(segments): logger.info( "Split long speaker segments by low energy: {} -> {}, max={}s", len(segments), len(split_segments), settings.MAX_SEGMENT_SEC, ) return split_segments def split_audio_by_speakers( self, audio_path: str, output_dir: Optional[str] = None, ) -> List[SpeakerSegment]: """完整的说话人分离流程 流程： 1. 执行CAM++说话人分离 2. 智能合并短片段（两层合并策略） - 第一层：<10s片段累积合并 - 第二层：60s累积合并 3. 提取音频数据，保存临时文件 Args: audio_path: 音频文件路径 output_dir: 输出目录 Returns: SpeakerSegment 列表 """ try: # 1. 执行说话人分离 raw_segments = self.diarize(audio_path) if not raw_segments: logger.warning("说话人分离未检测到任何片段") return [] # 2. 智能合并短片段（第一个<10s的同说话人片段向后合并） final_segments = self.merge_short_segments(raw_segments) # 3. Load audio before low-energy splitting and segment extraction. logger.info("加载音频并提取片段...") audio_data, sr = librosa.load(audio_path, sr=self.DEFAULT_SAMPLE_RATE) sample_rate = int(sr) final_segments = self.split_long_segments( final_segments, audio_data, sample_rate, ) logger.info(f"智能合并完成: {len(raw_segments)} → {len(final_segments)} 个片段") output_dir = output_dir or settings.TEMP_DIR os.makedirs(output_dir, exist_ok=True) for idx, seg in enumerate(final_segments): start_sample = int(seg.start_ms / 1000 * sample_rate) end_sample = int(seg.end_ms / 1000 * sample_rate) seg.audio_data = audio_data[start_sample:end_sample] # 保存临时文件 temp_file = tempfile.NamedTemporaryFile( delete=False, suffix=".wav", dir=output_dir, prefix=f"{seg.speaker_id}_{idx:03d}_", ) temp_path = temp_file.name temp_file.close() sf.write(temp_path, seg.audio_data, sample_rate) seg.temp_file = temp_path # 统计 unique_speakers = sorted(set(seg.speaker_id for seg in final_segments)) logger.info( f"音频分割完成: {len(final_segments)} 个片段, " f"{len(unique_speakers)} 个说话人" ) for spk in unique_speakers: spk_segs = [s for s in final_segments if s.speaker_id == spk] total_time = sum(s.duration_sec for s in spk_segs) logger.info(f" {spk}: {len(spk_segs)} 片段, {total_time:.2f}s") return final_segments except Exception as e: logger.error(f"说话人分离流程失败: {e}") raise DefaultServerErrorException(f"说话人分离失败: {str(e)}") @staticmethod def cleanup_segments(segments: List[SpeakerSegment]) -> None: """清理临时文件""" for seg in segments: if seg.temp_file and os.path.exists(seg.temp_file): try: os.remove(seg.temp_file) except Exception as e: logger.warning(f"清理临时文件失败: {seg.temp_file}, {e}")