wehub-resource-sync
436 lines
21 KiB
Python
436 lines
21 KiB
Python
import os
|
|
import pickle
|
|
import shutil
|
|
import sys
|
|
from io import BufferedReader
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import yaml
|
|
from loguru import logger
|
|
from sklearn.metrics.pairwise import cosine_similarity
|
|
from tqdm import tqdm
|
|
from yeaudio.audio import AudioSegment
|
|
|
|
from mvector.data_utils.featurizer import AudioFeaturizer
|
|
from mvector.infer_utils.speaker_diarization import SpeakerDiarization
|
|
from mvector.models import build_model
|
|
from mvector.utils.checkpoint import load_pretrained
|
|
from mvector.utils.utils import dict_to_object, print_arguments, convert_string_based_on_type
|
|
|
|
|
|
class MVectorPredictor:
|
|
def __init__(self,
|
|
configs,
|
|
threshold=0.6,
|
|
audio_db_path=None,
|
|
model_path='models/CAMPPlus_Fbank/best_model/',
|
|
use_gpu=True,
|
|
overwrites=None,
|
|
log_level="info"):
|
|
"""声纹识别预测工具
|
|
|
|
:param configs: 配置文件路径，或者模型名称，如果是模型名称则会使用默认的配置文件
|
|
:param threshold: 判断是否为同一个人的阈值
|
|
:param audio_db_path: 声纹库路径
|
|
:param model_path: 导出的预测模型文件夹路径
|
|
:param use_gpu: 是否使用GPU预测
|
|
:param overwrites: 覆盖配置文件中的参数，比如"train_conf.max_epoch=100"，多个用逗号隔开
|
|
:param log_level: 打印的日志等级，可选值有："debug", "info", "warning", "error"
|
|
"""
|
|
if use_gpu:
|
|
assert (torch.cuda.is_available()), 'GPU不可用'
|
|
self.device = torch.device("cuda")
|
|
else:
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
|
|
self.device = torch.device("cpu")
|
|
self.threshold = threshold
|
|
self.log_level = log_level.upper()
|
|
logger.remove()
|
|
logger.add(sink=sys.stdout, level=self.log_level)
|
|
# 读取配置文件
|
|
if isinstance(configs, str):
|
|
# 获取当前程序绝对路径
|
|
absolute_path = os.path.dirname(__file__)
|
|
# 获取默认配置文件路径
|
|
config_path = os.path.join(absolute_path, f"configs/{configs}.yml")
|
|
configs = config_path if os.path.exists(config_path) else configs
|
|
with open(configs, 'r', encoding='utf-8') as f:
|
|
configs = yaml.load(f.read(), Loader=yaml.FullLoader)
|
|
self.configs = dict_to_object(configs)
|
|
# 覆盖配置文件中的参数
|
|
if overwrites:
|
|
overwrites = overwrites.split(",")
|
|
for overwrite in overwrites:
|
|
keys, value = overwrite.strip().split("=")
|
|
attrs = keys.split('.')
|
|
current_level = self.configs
|
|
for attr in attrs[:-1]:
|
|
current_level = getattr(current_level, attr)
|
|
before_value = getattr(current_level, attrs[-1])
|
|
setattr(current_level, attrs[-1], convert_string_based_on_type(before_value, value))
|
|
# 打印配置信息
|
|
print_arguments(configs=self.configs)
|
|
self._audio_featurizer = AudioFeaturizer(feature_method=self.configs.preprocess_conf.feature_method,
|
|
use_hf_model=self.configs.preprocess_conf.get('use_hf_model', False),
|
|
method_args=self.configs.preprocess_conf.get('method_args', {}))
|
|
# 获取模型
|
|
backbone = build_model(input_size=self._audio_featurizer.feature_dim, configs=self.configs)
|
|
self.predictor = nn.Sequential(backbone)
|
|
self.predictor.to(self.device)
|
|
# 加载模型
|
|
if os.path.isdir(model_path):
|
|
model_path = os.path.join(model_path, 'model.pth')
|
|
assert os.path.exists(model_path), f"{model_path} 模型不存在！"
|
|
self.predictor = load_pretrained(self.predictor, model_path, use_gpu=use_gpu)
|
|
logger.info(f"成功加载模型参数：{model_path}")
|
|
self.predictor.eval()
|
|
|
|
# 声纹库的声纹特征
|
|
self.audio_feature = None
|
|
# 每个用户的平均声纹特征
|
|
self.audio_feature_mean = None
|
|
# 声纹特征对应的用户名
|
|
self.users_name = []
|
|
# 声纹特征对应的声纹文件路径
|
|
self.users_audio_path = []
|
|
# 每个用户的平均声纹特征对应的用户名称
|
|
self.users_name_mean = []
|
|
# 加载声纹库
|
|
self.audio_db_path = audio_db_path
|
|
if self.audio_db_path is not None:
|
|
self.audio_indexes_path = os.path.join(audio_db_path, "audio_indexes.bin")
|
|
# 加载声纹库中的声纹
|
|
self.__load_audio_db(self.audio_db_path)
|
|
# 说话人日志
|
|
self.speaker_diarize = SpeakerDiarization()
|
|
|
|
# 加载声纹特征索引
|
|
def __load_audio_indexes(self):
|
|
# 如果存在声纹特征索引文件就加载
|
|
if not os.path.exists(self.audio_indexes_path): return
|
|
with open(self.audio_indexes_path, "rb") as f:
|
|
indexes = pickle.load(f)
|
|
for name, feature, path in zip(indexes["users_name"], indexes["faces_feature"],
|
|
indexes["users_image_path"]):
|
|
if not os.path.exists(path): continue
|
|
self.users_name.append(name)
|
|
self.users_audio_path.append(path)
|
|
if self.audio_feature is None:
|
|
self.audio_feature = feature[np.newaxis, :]
|
|
else:
|
|
self.audio_feature = np.vstack((self.audio_feature, feature))
|
|
|
|
# 保存声纹特征索引
|
|
def __write_index(self):
|
|
with open(self.audio_indexes_path, "wb") as f:
|
|
pickle.dump({"users_name": self.users_name,
|
|
"faces_feature": self.audio_feature,
|
|
"users_image_path": self.users_audio_path}, f)
|
|
|
|
# 加载声纹库中的声纹
|
|
def __load_audio_db(self, audio_db_path):
|
|
# 先加载声纹特征索引
|
|
self.__load_audio_indexes()
|
|
os.makedirs(audio_db_path, exist_ok=True)
|
|
audios_path = []
|
|
for name in os.listdir(audio_db_path):
|
|
audio_dir = os.path.join(audio_db_path, name)
|
|
if not os.path.isdir(audio_dir): continue
|
|
for file in os.listdir(audio_dir):
|
|
audios_path.append(os.path.join(audio_dir, file).replace('\\', '/'))
|
|
# 声纹库没数据就跳过
|
|
if len(audios_path) == 0: return
|
|
logger.info('正在加载声纹库数据...')
|
|
input_audios = []
|
|
for audio_path in tqdm(audios_path, desc='加载声纹库数据'):
|
|
# 如果声纹特征已经在索引就跳过
|
|
if audio_path in self.users_audio_path: continue
|
|
# 读取声纹库音频
|
|
audio_segment = self._load_audio(audio_path)
|
|
# 获取用户名
|
|
user_name = os.path.basename(os.path.dirname(audio_path))
|
|
self.users_name.append(user_name)
|
|
self.users_audio_path.append(audio_path)
|
|
input_audios.append(audio_segment.samples)
|
|
# 处理一批数据
|
|
if len(input_audios) == self.configs.dataset_conf.eval_conf.batch_size:
|
|
features = self.predict_batch(input_audios)
|
|
if self.audio_feature is None:
|
|
self.audio_feature = features
|
|
else:
|
|
self.audio_feature = np.vstack((self.audio_feature, features))
|
|
input_audios = []
|
|
# 处理不满一批的数据
|
|
if len(input_audios) != 0:
|
|
features = self.predict_batch(input_audios)
|
|
if self.audio_feature is None:
|
|
self.audio_feature = features
|
|
else:
|
|
self.audio_feature = np.vstack((self.audio_feature, features))
|
|
assert len(self.audio_feature) == len(self.users_name) == len(self.users_audio_path), '加载的数量对不上！'
|
|
# 将声纹特征保存到索引文件中
|
|
self.__write_index()
|
|
# 计算平均特征，用于检索
|
|
for name in set(self.users_name):
|
|
indexes = [idx for idx, val in enumerate(self.users_name) if val == name]
|
|
feature = self.audio_feature[indexes].mean(axis=0)
|
|
if self.audio_feature_mean is None:
|
|
self.audio_feature_mean = feature[np.newaxis, :]
|
|
else:
|
|
self.audio_feature_mean = np.vstack((self.audio_feature_mean, feature))
|
|
self.users_name_mean.append(name)
|
|
if len(self.audio_feature_mean.shape) == 1:
|
|
self.audio_feature_mean = self.audio_feature_mean[np.newaxis, :]
|
|
logger.info(f'声纹库数据加载完成，一共有{len(self.audio_feature_mean)}个用户，分别是：{self.users_name_mean}')
|
|
|
|
# 特征进行归一化
|
|
@staticmethod
|
|
def normalize_features(features):
|
|
return features / np.linalg.norm(features, axis=1, keepdims=True)
|
|
|
|
# 声纹检索
|
|
def __retrieval(self, np_feature):
|
|
if isinstance(np_feature, list):
|
|
np_feature = np.array(np_feature)
|
|
labels = []
|
|
np_feature = self.normalize_features(np_feature.astype(np.float32))
|
|
if len(self.audio_feature_mean.shape) == 1:
|
|
self.audio_feature_mean = self.audio_feature_mean[np.newaxis, :]
|
|
similarities = cosine_similarity(np_feature, self.audio_feature_mean)
|
|
for sim in similarities:
|
|
idx = np.argmax(sim)
|
|
sim = sim[idx]
|
|
if sim >= self.threshold:
|
|
sim = round(float(sim), 5)
|
|
labels.append([self.users_name_mean[idx], sim])
|
|
else:
|
|
labels.append([None, None])
|
|
return labels
|
|
|
|
def _load_audio(self, audio_data, sample_rate=16000):
|
|
"""加载音频
|
|
:param audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy，AudioSegment对象。如果是字节的话，必须是完整的字节文件
|
|
:param sample_rate: 如果传入的事numpy数据，需要指定采样率
|
|
:return: 识别的文本结果和解码的得分数
|
|
"""
|
|
# 加载音频文件，并进行预处理
|
|
if isinstance(audio_data, str):
|
|
audio_segment = AudioSegment.from_file(audio_data)
|
|
elif isinstance(audio_data, BufferedReader):
|
|
audio_segment = AudioSegment.from_file(audio_data)
|
|
elif isinstance(audio_data, np.ndarray):
|
|
audio_segment = AudioSegment.from_ndarray(audio_data, sample_rate)
|
|
elif isinstance(audio_data, bytes):
|
|
audio_segment = AudioSegment.from_bytes(audio_data)
|
|
elif isinstance(audio_data, AudioSegment):
|
|
audio_segment = audio_data
|
|
else:
|
|
raise Exception(f'不支持该数据类型，当前数据类型为：{type(audio_data)}')
|
|
assert audio_segment.duration >= self.configs.dataset_conf.dataset.min_duration, \
|
|
f'音频太短，最小应该为{self.configs.dataset_conf.dataset.min_duration}s，当前音频为{audio_segment.duration}s'
|
|
# 重采样
|
|
if audio_segment.sample_rate != self.configs.dataset_conf.dataset.sample_rate:
|
|
audio_segment.resample(self.configs.dataset_conf.dataset.sample_rate)
|
|
# decibel normalization
|
|
if self.configs.dataset_conf.dataset.use_dB_normalization:
|
|
audio_segment.normalize(target_db=self.configs.dataset_conf.dataset.target_dB)
|
|
return audio_segment
|
|
|
|
def predict(self,
|
|
audio_data,
|
|
sample_rate=16000):
|
|
"""预测一个音频的特征
|
|
|
|
:param audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy，AudioSegment对象。如果是字节的话，必须是完整并带格式的字节文件
|
|
:param sample_rate: 如果传入的事numpy数据，需要指定采样率
|
|
:return: 声纹特征向量
|
|
"""
|
|
# 加载音频文件，并进行预处理
|
|
input_data = self._load_audio(audio_data=audio_data, sample_rate=sample_rate)
|
|
input_data = torch.tensor(input_data.samples, dtype=torch.float32).unsqueeze(0)
|
|
audio_feature = self._audio_featurizer(input_data).to(self.device)
|
|
# 执行预测
|
|
feature = self.predictor(audio_feature).data.cpu().numpy()[0]
|
|
return feature
|
|
|
|
def predict_batch(self, audios_data, sample_rate=16000, batch_size=32):
|
|
"""预测一批音频的特征
|
|
|
|
:param audios_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy，AudioSegment对象。如果是字节的话，必须是完整并带格式的字节文件
|
|
:param sample_rate: 如果传入的事numpy数据，需要指定采样率
|
|
:return: 声纹特征向量
|
|
"""
|
|
audios_data1 = []
|
|
for audio_data in audios_data:
|
|
# 加载音频文件，并进行预处理
|
|
input_data = self._load_audio(audio_data=audio_data, sample_rate=sample_rate)
|
|
audios_data1.append(input_data.samples)
|
|
# 找出音频长度最长的
|
|
batch = sorted(audios_data1, key=lambda a: a.shape[0], reverse=True)
|
|
max_audio_length = batch[0].shape[0]
|
|
input_size = len(batch)
|
|
# 以最大的长度创建0张量
|
|
inputs = np.zeros((input_size, max_audio_length), dtype=np.float32)
|
|
input_lens_ratio = []
|
|
for x in range(input_size):
|
|
tensor = audios_data1[x]
|
|
seq_length = tensor.shape[0]
|
|
# 将数据插入都0张量中，实现了padding
|
|
inputs[x, :seq_length] = tensor[:]
|
|
input_lens_ratio.append(seq_length / max_audio_length)
|
|
inputs = torch.tensor(inputs, dtype=torch.float32)
|
|
input_lens_ratio = torch.tensor(input_lens_ratio, dtype=torch.float32)
|
|
audio_feature = self._audio_featurizer(inputs, input_lens_ratio).to(self.device)
|
|
# 执行预测
|
|
features = []
|
|
for i in range(0, input_size, batch_size):
|
|
feature = self.predictor(audio_feature[i:i + batch_size]).data.cpu().numpy()
|
|
features.extend(feature)
|
|
features = np.array(features)
|
|
return features
|
|
|
|
def contrast(self, audio_data1, audio_data2):
|
|
"""声纹对比
|
|
|
|
param audio_data1: 需要对比的音频1，支持文件路径，文件对象，字节，numpy，AudioSegment对象。如果是字节的话，必须是完整的字节文件
|
|
param audio_data2: 需要对比的音频2，支持文件路径，文件对象，字节，numpy，AudioSegment对象。如果是字节的话，必须是完整的字节文件
|
|
|
|
return: 两个音频的相似度
|
|
"""
|
|
feature1 = self.predict(audio_data1)
|
|
feature2 = self.predict(audio_data2)
|
|
# 对角余弦值
|
|
dist = np.dot(feature1, feature2) / (np.linalg.norm(feature1) * np.linalg.norm(feature2))
|
|
return dist
|
|
|
|
def register(self,
|
|
audio_data,
|
|
user_name: str,
|
|
sample_rate=16000):
|
|
"""声纹注册
|
|
:param audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy。如果是字节的话，必须是完整的字节文件
|
|
:param user_name: 注册用户的名字
|
|
:param sample_rate: 如果传入的事numpy数据，需要指定采样率
|
|
:return: 识别的文本结果和解码的得分数
|
|
"""
|
|
# 加载音频文件
|
|
audio_segment = self._load_audio(audio_data=audio_data, sample_rate=sample_rate)
|
|
feature = self.predict(audio_data=audio_segment)
|
|
if self.audio_feature is None:
|
|
self.audio_feature = feature[np.newaxis, :]
|
|
else:
|
|
if len(self.audio_feature) == 0:
|
|
self.audio_feature = feature[np.newaxis, :]
|
|
else:
|
|
self.audio_feature = np.vstack((self.audio_feature, feature))
|
|
# 保存
|
|
if not os.path.exists(os.path.join(self.audio_db_path, user_name)):
|
|
audio_path = os.path.join(self.audio_db_path, user_name, '0.wav')
|
|
else:
|
|
audio_path = os.path.join(self.audio_db_path, user_name,
|
|
f'{len(os.listdir(os.path.join(self.audio_db_path, user_name)))}.wav')
|
|
os.makedirs(os.path.dirname(audio_path), exist_ok=True)
|
|
audio_segment.to_wav_file(audio_path)
|
|
self.users_audio_path.append(audio_path.replace('\\', '/'))
|
|
self.users_name.append(user_name)
|
|
self.__write_index()
|
|
# 更新检索的特征
|
|
if user_name in self.users_name_mean:
|
|
index = self.users_name_mean.index(user_name)
|
|
indexes = [idx for idx, val in enumerate(self.users_name) if val == user_name]
|
|
feature = self.audio_feature[indexes].mean(axis=0)
|
|
self.audio_feature_mean[index] = feature
|
|
else:
|
|
self.users_name_mean.append(user_name)
|
|
if self.audio_feature_mean is None:
|
|
self.audio_feature_mean = feature[np.newaxis, :]
|
|
else:
|
|
if len(self.audio_feature_mean) == 0:
|
|
self.audio_feature_mean = feature[np.newaxis, :]
|
|
else:
|
|
self.audio_feature_mean = np.vstack((self.audio_feature_mean, feature))
|
|
return True, "注册成功"
|
|
|
|
def recognition(self, audio_data, threshold=None, sample_rate=16000):
|
|
"""声纹识别
|
|
|
|
Args:
|
|
audio_data (str, file-like object, bytes, numpy.ndarray, AudioSegment): 需要识别的数据
|
|
threshold (float): 判断的阈值，如果为None则用创建对象时使用的阈值。默认为None。
|
|
sample_rate (int): 如果传入的是numpy数组，需要指定采样率。默认为16000。
|
|
Returns:
|
|
str: 识别的用户名称，如果为None，即没有识别到用户。
|
|
"""
|
|
if threshold:
|
|
self.threshold = threshold
|
|
if self.audio_feature_mean is None:
|
|
logger.warning("声纹库没有任何数据")
|
|
return [None, None]
|
|
feature = self.predict(audio_data, sample_rate=sample_rate)
|
|
result = self.__retrieval(np_feature=np.array([feature]))[0]
|
|
return result
|
|
|
|
def get_users(self):
|
|
"""获取所有用户
|
|
|
|
return: 所有用户
|
|
"""
|
|
return self.users_name
|
|
|
|
def remove_user(self, user_name):
|
|
"""删除用户
|
|
|
|
:param user_name: 用户名
|
|
:return:
|
|
"""
|
|
if user_name in self.users_name and user_name in self.users_name_mean:
|
|
indexes = [i for i in range(len(self.users_name)) if self.users_name[i] == user_name]
|
|
for index in sorted(indexes, reverse=True):
|
|
del self.users_name[index]
|
|
del self.users_audio_path[index]
|
|
self.audio_feature = np.delete(self.audio_feature, index, axis=0)
|
|
self.__write_index()
|
|
shutil.rmtree(os.path.join(self.audio_db_path, user_name))
|
|
# 删除检索内的特征
|
|
index = self.users_name_mean.index(user_name)
|
|
del self.users_name_mean[index]
|
|
self.audio_feature_mean = np.delete(self.audio_feature_mean, index, axis=0)
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def speaker_diarization(self, audio_data, sample_rate=16000, speaker_num=None, search_audio_db=False):
|
|
"""说话人日志识别
|
|
|
|
Args:
|
|
audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy。如果是字节的话，必须是完整并带格式的字节文件
|
|
sample_rate (int): 如果传入的是numpy数据，需要指定采样率
|
|
speaker_num (int): 预期的说话人数量，提供说话人数量可以提高准确率
|
|
search_audio_db (bool): 是否在数据库中搜索与输入音频最匹配的音频进行识别
|
|
Returns:
|
|
list: 说话人日志识别结果
|
|
"""
|
|
input_data = self._load_audio(audio_data=audio_data, sample_rate=sample_rate)
|
|
segments = self.speaker_diarize.segments_audio(input_data)
|
|
segments_data = [segment[2] for segment in segments]
|
|
features = self.predict_batch(segments_data, sample_rate=sample_rate)
|
|
labels, spk_center_embeddings = self.speaker_diarize.clustering(features, speaker_num=speaker_num)
|
|
outputs = self.speaker_diarize.postprocess(segments, labels)
|
|
if search_audio_db:
|
|
assert self.audio_feature is not None, "数据库中没有音频数据，请先指定说话人特征数据库或者注册说话人"
|
|
names = self.__retrieval(np_feature=spk_center_embeddings)
|
|
results = []
|
|
for output in outputs:
|
|
name = names[output['speaker']][0]
|
|
result = {
|
|
'speaker': name if name else f"陌生人{output['speaker']}",
|
|
'start': output['start'],
|
|
'end': output['end']
|
|
}
|
|
results.append(result)
|
|
outputs = results
|
|
return outputs
|
|
|