96 lines
2.0 KiB
Kotlin
96 lines
2.0 KiB
Kotlin
package com.k2fsa.sherpa.onnx
|
|
|
|
import android.content.res.AssetManager
|
|
import android.util.Log
|
|
|
|
private val TAG = "sherpa-onnx"
|
|
|
|
data class OfflineZipformerAudioTaggingModelConfig (
|
|
val model: String,
|
|
)
|
|
|
|
data class AudioTaggingModelConfig (
|
|
var zipformer: OfflineZipformerAudioTaggingModelConfig,
|
|
var numThreads: Int = 1,
|
|
var debug: Boolean = false,
|
|
var provider: String = "cpu",
|
|
)
|
|
|
|
data class AudioTaggingConfig (
|
|
var model: AudioTaggingModelConfig,
|
|
var labels: String,
|
|
var topK: Int = 5,
|
|
)
|
|
|
|
data class AudioEvent (
|
|
val name: String,
|
|
val index: Int,
|
|
val prob: Float,
|
|
)
|
|
|
|
class AudioTagging(
|
|
assetManager: AssetManager? = null,
|
|
config: AudioTaggingConfig,
|
|
) {
|
|
private var ptr: Long
|
|
|
|
init {
|
|
ptr = if (assetManager != null) {
|
|
newFromAsset(assetManager, config)
|
|
} else {
|
|
newFromFile(config)
|
|
}
|
|
}
|
|
|
|
protected fun finalize() {
|
|
if(ptr != 0L) {
|
|
delete(ptr)
|
|
ptr = 0
|
|
}
|
|
}
|
|
|
|
fun release() = finalize()
|
|
|
|
fun createStream(): OfflineStream {
|
|
val p = createStream(ptr)
|
|
return OfflineStream(p)
|
|
}
|
|
|
|
fun compute(stream: OfflineStream, topK: Int=-1): ArrayList<AudioEvent> {
|
|
var events :Array<Any> = compute(ptr, stream.ptr, topK)
|
|
val ans = ArrayList<AudioEvent>()
|
|
|
|
for (e in events) {
|
|
val p :Array<Any> = e as Array<Any>
|
|
ans.add(AudioEvent(
|
|
name=p[0] as String,
|
|
index=p[1] as Int,
|
|
prob=p[2] as Float,
|
|
))
|
|
}
|
|
|
|
return ans
|
|
}
|
|
|
|
private external fun newFromAsset(
|
|
assetManager: AssetManager,
|
|
config: AudioTaggingConfig,
|
|
): Long
|
|
|
|
private external fun newFromFile(
|
|
config: AudioTaggingConfig,
|
|
): Long
|
|
|
|
private external fun delete(ptr: Long)
|
|
|
|
private external fun createStream(ptr: Long): Long
|
|
|
|
private external fun compute(ptr: Long, streamPtr: Long, topK: Int): Array<Any>
|
|
|
|
companion object {
|
|
init {
|
|
System.loadLibrary("sherpa-onnx-jni")
|
|
}
|
|
}
|
|
}
|