335 lines
12 KiB
Kotlin
335 lines
12 KiB
Kotlin
package com.k2fsa.sherpa.onnx
|
|
|
|
import android.content.res.AssetManager
|
|
import android.util.Log
|
|
import kotlin.math.sqrt
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Native speaker-embedding extractor (wraps libsherpa-onnx-jni).
|
|
// Ported from android/SherpaOnnxSpeakerIdentification so both apps stay
|
|
// compatible with the same JNI symbols.
|
|
// ---------------------------------------------------------------------------
|
|
class SpeakerEmbeddingExtractor(
|
|
assetManager: AssetManager? = null,
|
|
config: SpeakerEmbeddingExtractorConfig,
|
|
) {
|
|
private var ptr: Long
|
|
|
|
init {
|
|
ptr = if (assetManager != null) {
|
|
newFromAsset(assetManager, config)
|
|
} else {
|
|
newFromFile(config)
|
|
}
|
|
require(ptr != 0L) {
|
|
"Invalid SpeakerEmbeddingExtractorConfig: failed to create native SpeakerEmbeddingExtractor"
|
|
}
|
|
}
|
|
|
|
protected fun finalize() {
|
|
if (ptr != 0L) {
|
|
delete(ptr)
|
|
ptr = 0
|
|
}
|
|
}
|
|
|
|
fun release() = finalize()
|
|
|
|
fun createStream(): OnlineStream {
|
|
val p = createStream(ptr)
|
|
return OnlineStream(p)
|
|
}
|
|
|
|
fun isReady(stream: OnlineStream) = isReady(ptr, stream.ptr)
|
|
fun compute(stream: OnlineStream) = compute(ptr, stream.ptr)
|
|
fun dim() = dim(ptr)
|
|
|
|
private external fun newFromAsset(
|
|
assetManager: AssetManager,
|
|
config: SpeakerEmbeddingExtractorConfig,
|
|
): Long
|
|
|
|
private external fun newFromFile(config: SpeakerEmbeddingExtractorConfig): Long
|
|
private external fun delete(ptr: Long)
|
|
private external fun createStream(ptr: Long): Long
|
|
private external fun isReady(ptr: Long, streamPtr: Long): Boolean
|
|
private external fun compute(ptr: Long, streamPtr: Long): FloatArray
|
|
private external fun dim(ptr: Long): Int
|
|
|
|
companion object {
|
|
init {
|
|
System.loadLibrary("sherpa-onnx-jni")
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// SpeakerRegistry: lightweight in-Kotlin speaker store + cosine matching.
|
|
//
|
|
// We avoid the native SpeakerEmbeddingManager on purpose: it cannot rename a
|
|
// speaker once enrolled, which we need for the "auto-label then assign a real
|
|
// name" flow. Matching here prefers named entries over auto (S0/S1/...) ones.
|
|
// ---------------------------------------------------------------------------
|
|
class SpeakerRegistry {
|
|
private val named = ArrayList<Pair<String, FloatArray>>()
|
|
private val auto = ArrayList<Pair<String, FloatArray>>()
|
|
|
|
val namedNames: List<String>
|
|
get() = named.map { it.first }
|
|
|
|
// The rolling window checker runs on its own thread while the endpoint
|
|
// identification runs on the recording thread: serialize everything.
|
|
@Synchronized
|
|
fun clear() {
|
|
named.clear()
|
|
auto.clear()
|
|
}
|
|
|
|
// Returns the best label for embedding `e`, or null if no speaker is close
|
|
// enough (below threshold). Registered names win ties against auto labels.
|
|
@Synchronized
|
|
fun search(e: FloatArray, threshold: Float): String? {
|
|
match(named, e, threshold)?.let { return it }
|
|
match(auto, e, threshold)?.let { return it }
|
|
return null
|
|
}
|
|
|
|
@Synchronized
|
|
fun addAuto(name: String, e: FloatArray) {
|
|
auto.add(Pair(name, l2normalize(e)))
|
|
}
|
|
|
|
// Enroll (or overwrite) a real name for a given embedding.
|
|
@Synchronized
|
|
fun addNamed(name: String, e: FloatArray) {
|
|
val v = l2normalize(e)
|
|
named.removeAll { it.first == name }
|
|
// Drop any auto entry whose embedding now equals this named one so the
|
|
// same voice is not also reported as an anonymous Sx id.
|
|
named.add(Pair(name, v))
|
|
}
|
|
|
|
/**
|
|
* Best cosine similarity of `e` against all enrolled speakers (0..1 scale,
|
|
* may be slightly noisy). Used by the rolling checker to tell "no idea
|
|
* whose voice this is" from "definitely the current speaker".
|
|
*/
|
|
@Synchronized
|
|
fun bestScore(e: FloatArray): Float {
|
|
val q = l2normalize(e)
|
|
var best = -1.0f
|
|
for ((_, emb) in named) best = maxOf(best, dot(q, emb))
|
|
for ((_, emb) in auto) best = maxOf(best, dot(q, emb))
|
|
return best
|
|
}
|
|
|
|
/**
|
|
* 最近邻兜底(收紧版,P1):按"名字级"聚合所有嵌入,返回最优名与
|
|
* 名级最高/次高分。同人多嵌入只算一个竞争者,优势差判定才有意义;
|
|
* 调用方用 score≥minScore 且 (score-secondScore)≥margin 才归入,
|
|
* 模糊时判新人(宁多勿错)。
|
|
*/
|
|
class Rank(val name: String?, val score: Float, val secondScore: Float)
|
|
|
|
@Synchronized
|
|
fun nearestRank(e: FloatArray): Rank {
|
|
val q = l2normalize(e)
|
|
val byName = LinkedHashMap<String, Float>()
|
|
for ((name, emb) in named) {
|
|
val s = dot(q, emb)
|
|
if (s > (byName[name] ?: -2f)) byName[name] = s
|
|
}
|
|
for ((name, emb) in auto) {
|
|
val s = dot(q, emb)
|
|
if (s > (byName[name] ?: -2f)) byName[name] = s
|
|
}
|
|
var best: String? = null
|
|
var bestScore = -2f
|
|
var secondScore = -2f
|
|
for ((name, s) in byName) {
|
|
if (s > bestScore) { secondScore = bestScore; bestScore = s; best = name }
|
|
else if (s > secondScore) secondScore = s
|
|
}
|
|
return Rank(best, bestScore, secondScore)
|
|
}
|
|
|
|
/**
|
|
* 点击标签命名用:把内部键 S{n}(自动注册、未命名)升级为真名,
|
|
* 或把已有真名改成另一个真名。同名允许多条嵌入并存 —— 之前误判成
|
|
* 两个说话人、后被命名为同一名字的两组声纹,自此 search 都返回同一名。
|
|
*/
|
|
@Synchronized
|
|
fun renameToName(oldKey: String, name: String): Boolean {
|
|
val ai = auto.indexOfFirst { it.first == oldKey }
|
|
if (ai >= 0) {
|
|
named.add(Pair(name, auto[ai].second))
|
|
auto.removeAt(ai)
|
|
return true
|
|
}
|
|
var changed = false
|
|
for (i in named.indices) {
|
|
if (named[i].first == oldKey) {
|
|
named[i] = Pair(name, named[i].second)
|
|
changed = true
|
|
}
|
|
}
|
|
return changed
|
|
}
|
|
|
|
private fun match(
|
|
list: List<Pair<String, FloatArray>>,
|
|
e: FloatArray,
|
|
threshold: Float,
|
|
): String? {
|
|
if (list.isEmpty()) return null
|
|
val q = l2normalize(e)
|
|
var bestName: String? = null
|
|
var bestScore = -1.0f
|
|
for ((name, emb) in list) {
|
|
val s = dot(q, emb)
|
|
if (s > bestScore) {
|
|
bestScore = s
|
|
bestName = name
|
|
}
|
|
}
|
|
return if (bestScore >= threshold) bestName else null
|
|
}
|
|
|
|
private fun l2normalize(v: FloatArray): FloatArray {
|
|
var sum = 0.0f
|
|
for (x in v) sum += x * x
|
|
if (sum <= 0f) return v.copyOf()
|
|
val inv = 1.0f / sqrt(sum)
|
|
return FloatArray(v.size) { v[it] * inv }
|
|
}
|
|
|
|
private fun dot(a: FloatArray, b: FloatArray): Float {
|
|
var s = 0.0f
|
|
val n = minOf(a.size, b.size)
|
|
for (i in 0 until n) s += a[i] * b[i]
|
|
return s
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Singleton holder shared by the UI thread, recording thread and the rolling
|
|
// speaker-check thread.
|
|
//
|
|
// 双模型方案(E):准确模型(eres2netv2)负责 endpoint 权威判定与注册;
|
|
// 轻量模型(eres2net base,可选)只负责滚动换人检测(推理快 ~4 倍)。
|
|
// ⚠ 两个模型的 embedding 不互通:各自配一份 registry,注册时双写镜像。
|
|
// ---------------------------------------------------------------------------
|
|
object SpeakerRecognition {
|
|
@Volatile private var _extractor: SpeakerEmbeddingExtractor? = null
|
|
@Volatile private var _fastExtractor: SpeakerEmbeddingExtractor? = null
|
|
val extractor: SpeakerEmbeddingExtractor
|
|
get() = _extractor!!
|
|
|
|
/** accurate-space registry (endpoint verdict / naming source of truth) */
|
|
val registry = SpeakerRegistry()
|
|
/** fast-space mirror (rolling-check search only; may stay empty) */
|
|
val fastRegistry = SpeakerRegistry()
|
|
|
|
var autoCounter: Int = 0
|
|
private set
|
|
|
|
private val accLock = Any()
|
|
private val fastLock = Any()
|
|
|
|
// autoCounter is touched from worker threads while resetSpeakers runs on
|
|
// the UI thread — keep them serialized.
|
|
@Synchronized
|
|
fun nextAutoLabel(): String = "S${autoCounter++}"
|
|
|
|
@Synchronized
|
|
fun resetSpeakers() {
|
|
registry.clear()
|
|
fastRegistry.clear()
|
|
autoCounter = 0
|
|
}
|
|
|
|
// assetManager = null -> `model` is treated as an absolute file system path
|
|
fun initExtractor(assetManager: AssetManager?, model: String) {
|
|
synchronized(accLock) {
|
|
if (_extractor != null) return
|
|
Log.i(TAG, "Initializing speaker embedding extractor with $model")
|
|
_extractor = SpeakerEmbeddingExtractor(
|
|
assetManager = assetManager,
|
|
config = SpeakerEmbeddingExtractorConfig(
|
|
model = model,
|
|
numThreads = 2,
|
|
debug = false,
|
|
provider = "cpu",
|
|
)
|
|
)
|
|
}
|
|
}
|
|
|
|
// Optional lightweight model for the rolling speaker-change detector.
|
|
fun initFastExtractor(assetManager: AssetManager?, model: String) {
|
|
synchronized(fastLock) {
|
|
if (_fastExtractor != null) return
|
|
Log.i(TAG, "Initializing FAST speaker embedding extractor with $model")
|
|
_fastExtractor = SpeakerEmbeddingExtractor(
|
|
assetManager = assetManager,
|
|
config = SpeakerEmbeddingExtractorConfig(
|
|
model = model,
|
|
numThreads = 1,
|
|
debug = false,
|
|
provider = "cpu",
|
|
)
|
|
)
|
|
}
|
|
}
|
|
|
|
fun fastReady(): Boolean = _fastExtractor != null
|
|
|
|
// Compute one embedding with the accurate extractor (endpoint verdicts).
|
|
fun computeEmbedding(samples: FloatArray, sampleRate: Int): FloatArray? =
|
|
synchronized(accLock) {
|
|
val ext = _extractor ?: return null
|
|
runCompute(ext, samples, sampleRate)
|
|
}
|
|
|
|
// Rolling-check embedding: fast space when available, else accurate space
|
|
// (in the latter case both registries see the same space anyway).
|
|
fun computeEmbeddingFast(samples: FloatArray, sampleRate: Int): FloatArray? {
|
|
val fast = _fastExtractor
|
|
return if (fast == null) computeEmbedding(samples, sampleRate)
|
|
else synchronized(fastLock) { runCompute(fast, samples, sampleRate) }
|
|
}
|
|
|
|
private fun runCompute(
|
|
ext: SpeakerEmbeddingExtractor,
|
|
samples: FloatArray,
|
|
sampleRate: Int,
|
|
): FloatArray? {
|
|
if (samples.isEmpty()) return null
|
|
val stream = ext.createStream()
|
|
stream.acceptWaveform(samples = samples, sampleRate = sampleRate)
|
|
stream.inputFinished()
|
|
val emb = if (ext.isReady(stream)) ext.compute(stream) else null
|
|
stream.release()
|
|
return emb
|
|
}
|
|
|
|
// Cosine similarity between two embeddings (L2-normalized internally),
|
|
// in [-1, 1]. Pure function, safe from any thread.
|
|
fun similarity(a: FloatArray, b: FloatArray): Float {
|
|
val n = minOf(a.size, b.size)
|
|
if (n == 0) return 0f
|
|
var dot = 0.0
|
|
var na = 0.0
|
|
var nb = 0.0
|
|
for (i in 0 until n) {
|
|
dot += a[i].toDouble() * b[i]
|
|
na += a[i].toDouble() * a[i]
|
|
nb += b[i].toDouble() * b[i]
|
|
}
|
|
if (na <= 0.0 || nb <= 0.0) return 0f
|
|
return (dot / (kotlin.math.sqrt(na) * kotlin.math.sqrt(nb))).toFloat()
|
|
}
|
|
|
|
private const val TAG = "sherpa-onnx-speaker"
|
|
}
|