64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
"""ECAPA-TDNN Voiceprint Engine — loads model, processes extract/verify tasks."""
|
|
import torch
|
|
from longtasks.longtasks import LongTasks
|
|
from appPublic.worker import awaitify
|
|
from appPublic.jsonConfig import getConfig
|
|
from appPublic.log import debug
|
|
from speechbrain.inference.speaker import SpeakerRecognition
|
|
import numpy as np
|
|
|
|
|
|
class VoiceprintEngine(LongTasks):
|
|
def __init__(self):
|
|
self.config = getConfig()
|
|
super().__init__(self.config.redis_url, 'voiceprint', worker_cnt=self.config.worker_cnt)
|
|
self.load_model()
|
|
|
|
def load_model(self):
|
|
device = self.config.device
|
|
debug(f'loading ECAPA-TDNN on {device}...')
|
|
self.verifier = SpeakerRecognition.from_hparams(
|
|
source="speechbrain/spkrec-ecapa-voxceleb",
|
|
savedir="/share/models/ecapa-tdnn",
|
|
run_opts={"device": device}
|
|
)
|
|
debug('ECAPA-TDNN loaded')
|
|
|
|
async def process_task(self, payload, workerid=None):
|
|
task_type = payload.get('task_type', 'extract')
|
|
audio_file = payload.get('audio_file', '')
|
|
if not audio_file:
|
|
return {'status': 'FAILED', 'result': 'missing audio_file'}
|
|
|
|
if task_type == 'extract':
|
|
f = awaitify(self._extract)
|
|
return await f(audio_file)
|
|
elif task_type == 'verify':
|
|
ref_file = payload.get('reference_file', '')
|
|
if not ref_file:
|
|
return {'status': 'FAILED', 'result': 'missing reference_file'}
|
|
f = awaitify(self._verify)
|
|
return await f(audio_file, ref_file)
|
|
return {'status': 'FAILED', 'result': f'unknown task_type: {task_type}'}
|
|
|
|
def _extract(self, audio_path):
|
|
signal = self.verifier.load_audio(audio_path, 16000)
|
|
t = torch.tensor(signal).unsqueeze(0).to(self.config.device)
|
|
emb = self.verifier.encode_batch(t)
|
|
vec = emb.squeeze().cpu().numpy().tolist()
|
|
return {
|
|
'status': 'SUCCEEDED',
|
|
'embedding': vec,
|
|
'embedding_dim': len(vec),
|
|
'usage': {'audio_duration': round(len(signal) / 16000, 2)},
|
|
}
|
|
|
|
def _verify(self, audio_path, ref_path):
|
|
score, pred = self.verifier.verify_files(audio_path, ref_path)
|
|
return {
|
|
'status': 'SUCCEEDED',
|
|
'similarity': round(float(score), 4),
|
|
'is_same_speaker': bool(pred),
|
|
'usage': {},
|
|
}
|