diff --git a/app/pipeline/recognition.py b/app/pipeline/recognition.py new file mode 100644 index 0000000..d017d3d --- /dev/null +++ b/app/pipeline/recognition.py @@ -0,0 +1,76 @@ +import struct +import io +import numpy as np + +# These are patched in tests; imported lazily at runtime to avoid docker-only dep +try: + from resemblyzer import VoiceEncoder, preprocess_wav +except ImportError: + VoiceEncoder = None + preprocess_wav = None + +_encoder = None + + +def get_encoder(): + global _encoder + if _encoder is None: + if VoiceEncoder is None: + raise ImportError("resemblyzer not installed") + _encoder = VoiceEncoder() + return _encoder + + +def embed_audio(wav_bytes: bytes) -> np.ndarray: + """Extract 256-dim d-vector from raw WAV bytes.""" + wav = preprocess_wav(io.BytesIO(wav_bytes)) + return get_encoder().embed_utterance(wav) + + +def pack_embedding(embedding: np.ndarray) -> bytes: + n = len(embedding) + return struct.pack(f"{n}f", *embedding) + + +def unpack_embedding(blob: bytes) -> np.ndarray: + n = len(blob) // 4 + return np.array(struct.unpack(f"{n}f", blob), dtype=np.float32) + + +def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + norm_a = np.linalg.norm(a) + norm_b = np.linalg.norm(b) + if norm_a == 0 or norm_b == 0: + return 0.0 + return float(np.dot(a, b) / (norm_a * norm_b)) + + +def find_best_speaker_match( + utterance_embedding: np.ndarray, + stored_embeddings: list, + known_threshold: float = 0.85, + ambiguous_threshold: float = 0.60, +) -> tuple: + """ + Returns (match_status, speaker_id, confidence). + Groups by speaker_id, takes MAX cosine similarity per speaker. + """ + if not stored_embeddings: + return "unknown", None, None + + best_by_speaker: dict = {} + for speaker_id, blob in stored_embeddings: + ref = unpack_embedding(blob) + sim = cosine_similarity(utterance_embedding, ref) + if speaker_id not in best_by_speaker or sim > best_by_speaker[speaker_id]: + best_by_speaker[speaker_id] = sim + + best_id = max(best_by_speaker, key=best_by_speaker.__getitem__) + best_sim = best_by_speaker[best_id] + + if best_sim >= known_threshold: + return "known", best_id, best_sim + elif best_sim >= ambiguous_threshold: + return "ambiguous", best_id, best_sim + else: + return "unknown", None, best_sim diff --git a/tests/test_recognition.py b/tests/test_recognition.py new file mode 100644 index 0000000..a7fcd0d --- /dev/null +++ b/tests/test_recognition.py @@ -0,0 +1,106 @@ +import numpy as np +import struct +import pytest +from unittest.mock import patch, MagicMock + + +def test_pack_unpack_roundtrip(): + from app.pipeline.recognition import pack_embedding, unpack_embedding + original = np.array([0.1, 0.2, 0.3, -0.5], dtype=np.float32) + packed = pack_embedding(original) + recovered = unpack_embedding(packed) + np.testing.assert_allclose(recovered, original, rtol=1e-6) + + +def test_cosine_similarity_identical(): + from app.pipeline.recognition import cosine_similarity + v = np.array([1.0, 0.0, 0.0], dtype=np.float32) + assert cosine_similarity(v, v) == pytest.approx(1.0) + + +def test_cosine_similarity_orthogonal(): + from app.pipeline.recognition import cosine_similarity + a = np.array([1.0, 0.0], dtype=np.float32) + b = np.array([0.0, 1.0], dtype=np.float32) + assert cosine_similarity(a, b) == pytest.approx(0.0) + + +def test_cosine_similarity_zero_vector(): + from app.pipeline.recognition import cosine_similarity + a = np.array([0.0, 0.0], dtype=np.float32) + b = np.array([1.0, 0.0], dtype=np.float32) + assert cosine_similarity(a, b) == 0.0 + + +def test_find_best_match_known(): + from app.pipeline.recognition import find_best_speaker_match, pack_embedding + ref = np.array([1.0, 0.0, 0.0], dtype=np.float32) + query = np.array([0.98, 0.2, 0.0], dtype=np.float32) + query /= np.linalg.norm(query) + ref /= np.linalg.norm(ref) + status, sid, conf = find_best_speaker_match( + query, [(1, pack_embedding(ref))], known_threshold=0.85, ambiguous_threshold=0.60 + ) + assert status == "known" + assert sid == 1 + assert conf >= 0.85 + + +def test_find_best_match_ambiguous(): + from app.pipeline.recognition import find_best_speaker_match, pack_embedding + ref = np.array([1.0, 0.0, 0.0], dtype=np.float32) + # cos similarity ~0.707 (45 degrees) + query = np.array([1.0, 1.0, 0.0], dtype=np.float32) + query /= np.linalg.norm(query) + ref /= np.linalg.norm(ref) + status, sid, conf = find_best_speaker_match( + query, [(1, pack_embedding(ref))], known_threshold=0.85, ambiguous_threshold=0.60 + ) + assert status == "ambiguous" + assert sid == 1 + + +def test_find_best_match_unknown(): + from app.pipeline.recognition import find_best_speaker_match, pack_embedding + ref = np.array([1.0, 0.0, 0.0], dtype=np.float32) + query = np.array([0.0, 1.0, 0.0], dtype=np.float32) # orthogonal = 0 similarity + status, sid, conf = find_best_speaker_match( + query, [(1, pack_embedding(ref))], known_threshold=0.85, ambiguous_threshold=0.60 + ) + assert status == "unknown" + assert sid is None + + +def test_find_best_match_empty_embeddings(): + from app.pipeline.recognition import find_best_speaker_match + query = np.array([1.0, 0.0], dtype=np.float32) + status, sid, conf = find_best_speaker_match(query, []) + assert status == "unknown" + assert sid is None + + +def test_find_best_match_takes_max_per_speaker(): + from app.pipeline.recognition import find_best_speaker_match, pack_embedding + # Speaker 1 has two embeddings — one poor, one good + ref_bad = np.array([0.0, 1.0, 0.0], dtype=np.float32) + ref_good = np.array([1.0, 0.0, 0.0], dtype=np.float32) + query = np.array([1.0, 0.0, 0.0], dtype=np.float32) + stored = [(1, pack_embedding(ref_bad)), (1, pack_embedding(ref_good))] + status, sid, conf = find_best_speaker_match( + query, stored, known_threshold=0.85, ambiguous_threshold=0.60 + ) + assert status == "known" + assert sid == 1 + + +def test_embed_audio_calls_resemblyzer(mocker): + from app.pipeline.recognition import embed_audio + mock_wav = np.zeros(16000, dtype=np.float32) + mock_emb = np.ones(256, dtype=np.float32) + mocker.patch("app.pipeline.recognition.preprocess_wav", return_value=mock_wav) + mock_encoder = MagicMock() + mock_encoder.embed_utterance.return_value = mock_emb + mocker.patch("app.pipeline.recognition.get_encoder", return_value=mock_encoder) + result = embed_audio(b"\x00" * 100) + assert result.shape == (256,) + mock_encoder.embed_utterance.assert_called_once_with(mock_wav)