189 lines
8.7 KiB
Python
189 lines
8.7 KiB
Python
import os
|
||
import pickle
|
||
import shutil
|
||
from io import BufferedReader
|
||
|
||
import numpy as np
|
||
import torch
|
||
import yaml
|
||
from sklearn.metrics.pairwise import cosine_similarity
|
||
from tqdm import tqdm
|
||
|
||
from mvector import SUPPORT_MODEL
|
||
from mvector.data_utils.audio import AudioSegment
|
||
from mvector.data_utils.featurizer import AudioFeaturizer
|
||
from mvector.models.ecapa_tdnn import EcapaTdnn
|
||
from mvector.models.fc import SpeakerIdetification
|
||
from mvector.utils.logger import setup_logger
|
||
from mvector.utils.utils import dict_to_object, print_arguments
|
||
|
||
logger = setup_logger(__name__)
|
||
|
||
|
||
class MVectorPredictor:
|
||
def __init__(self,
|
||
configs,
|
||
threshold=0.6,
|
||
model_path='models/ecapa_tdnn_FBank/best_model/',
|
||
use_gpu=True):
|
||
"""
|
||
声纹识别预测工具
|
||
:param configs: 配置参数
|
||
:param threshold: 判断是否为同一个人的阈值
|
||
:param model_path: 导出的预测模型文件夹路径
|
||
:param use_gpu: 是否使用GPU预测
|
||
"""
|
||
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.cdd_num = 5
|
||
self.threshold = threshold
|
||
# 读取配置文件
|
||
if isinstance(configs, str):
|
||
with open(configs, 'r', encoding='utf-8') as f:
|
||
configs = yaml.load(f.read(), Loader=yaml.FullLoader)
|
||
# print_arguments(configs=configs)
|
||
self.configs = dict_to_object(configs)
|
||
assert 'max_duration' in self.configs.dataset_conf, \
|
||
'【警告】，您貌似使用了旧的配置文件，如果你同时使用了旧的模型，这是错误的，请重新下载或者重新训练，否则只能回滚代码。'
|
||
assert self.configs.use_model in SUPPORT_MODEL, f'没有该模型：{self.configs.use_model}'
|
||
self._audio_featurizer = AudioFeaturizer(feature_conf=self.configs.feature_conf, **self.configs.preprocess_conf)
|
||
self._audio_featurizer.to(self.device)
|
||
# 获取模型
|
||
if self.configs.use_model == 'EcapaTdnn' or self.configs.use_model == 'ecapa_tdnn':
|
||
backbone = EcapaTdnn(input_size=self._audio_featurizer.feature_dim, **self.configs.model_conf)
|
||
else:
|
||
raise Exception(f'{self.configs.use_model} 模型不存在！')
|
||
model = SpeakerIdetification(backbone=backbone, num_class=self.configs.dataset_conf.num_speakers)
|
||
model.to(self.device)
|
||
# 加载模型
|
||
if os.path.isdir(model_path):
|
||
model_path = os.path.join(model_path, 'model.pt')
|
||
assert os.path.exists(model_path), f"{model_path} 模型不存在！"
|
||
if torch.cuda.is_available() and use_gpu:
|
||
model_state_dict = torch.load(model_path)
|
||
else:
|
||
model_state_dict = torch.load(model_path, map_location='cpu')
|
||
# 加载模型参数
|
||
model.load_state_dict(model_state_dict)
|
||
print(f"成功加载模型参数：{model_path}")
|
||
# 设置为评估模式
|
||
model.eval()
|
||
self.predictor = model.backbone
|
||
# 声纹库的声纹特征
|
||
self.audio_feature = None
|
||
|
||
def _load_audio(self, audio_data, sample_rate=16000):
|
||
"""加载音频
|
||
:param audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy。如果是字节的话，必须是完整的字节文件
|
||
: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)
|
||
else:
|
||
raise Exception(f'不支持该数据类型，当前数据类型为：{type(audio_data)}')
|
||
assert audio_segment.duration >= self.configs.dataset_conf.min_duration, \
|
||
f'音频太短，最小应该为{self.configs.dataset_conf.min_duration}s，当前音频为{audio_segment.duration}s'
|
||
# 重采样
|
||
if audio_segment.sample_rate != self.configs.dataset_conf.sample_rate:
|
||
audio_segment.resample(self.configs.dataset_conf.sample_rate)
|
||
# decibel normalization
|
||
if self.configs.dataset_conf.use_dB_normalization:
|
||
audio_segment.normalize(target_db=self.configs.dataset_conf.target_dB)
|
||
return audio_segment
|
||
|
||
def predict(self,
|
||
audio_data,
|
||
sample_rate=16000):
|
||
"""预测一个音频的特征
|
||
|
||
:param audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy。如果是字节的话，必须是完整并带格式的字节文件
|
||
: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, device=self.device).unsqueeze(0)
|
||
input_len_ratio = torch.tensor([1], dtype=torch.float32, device=self.device)
|
||
audio_feature, _ = self._audio_featurizer(input_data, input_len_ratio)
|
||
# 执行预测
|
||
feature = self.predictor(audio_feature).data.cpu().numpy()[0]
|
||
return feature
|
||
|
||
def predict_batch(self, audios_data, sample_rate=16000):
|
||
"""预测一批音频的特征
|
||
|
||
:param audios_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy。如果是字节的话，必须是完整并带格式的字节文件
|
||
: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]
|
||
batch_size = len(batch)
|
||
# 以最大的长度创建0张量
|
||
inputs = np.zeros((batch_size, max_audio_length), dtype='float32')
|
||
input_lens_ratio = []
|
||
for x in range(batch_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)
|
||
audios_data = torch.tensor(inputs, dtype=torch.float32, device=self.device)
|
||
input_lens_ratio = torch.tensor(input_lens_ratio, dtype=torch.float32, device=self.device)
|
||
audio_feature, _ = self._audio_featurizer(audios_data, input_lens_ratio)
|
||
# 执行预测
|
||
features = self.predictor(audio_feature).data.cpu().numpy()
|
||
return features
|
||
|
||
# 声纹对比
|
||
def contrast(self, audio_data1, audio_data2):
|
||
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 recognition(self, audio_data, threshold=None, sample_rate=16000):
|
||
"""声纹识别
|
||
:param audio_data: 需要识别的数据，支持文件路径，文件对象，字节，numpy。如果是字节的话，必须是完整的字节文件
|
||
:param threshold: 判断的阈值，如果为None则用创建对象时使用的阈值
|
||
:param sample_rate: 如果传入的事numpy数据，需要指定采样率
|
||
:return: 识别的用户名称，如果为None，即没有识别到用户
|
||
"""
|
||
if threshold:
|
||
self.threshold = threshold
|
||
feature = self.predict(audio_data, sample_rate=sample_rate)
|
||
name = self.__retrieval(np_feature=[feature])[0]
|
||
return name
|
||
|
||
def compare(self, feature1, feature2):
|
||
"""声纹对比
|
||
|
||
:param feature1: 特征1
|
||
:param feature2: 特征2
|
||
:return:
|
||
"""
|
||
# 对角余弦值
|
||
dist = np.dot(feature1, feature2) / (np.linalg.norm(feature1) * np.linalg.norm(feature2))
|
||
return dist
|
||
|