app/src/main/kotlin/com/honjimaku/subrep/Segmenter.kt (4254 bytes)
1 package com.honjimaku.subrep 2 3 import kotlin.math.sqrt 4 5 /** 6 * Cuts the sound into the pieces that go to Whisper: speech, closed by a pause. 7 * 8 * The rules are those of the desktop app. The sound comes in blocks of [blockMs]. A block is 9 * speech when its level is above the gate: the noise floor times [ratio], and at least 10 * [floor]. A piece starts [prerollMs] before the first speech block, and closes after 11 * [pauseMs] of silence, or at [maxMs] when the speech goes on. A piece with less than 12 * [minSpeechMs] of speech is dropped: a click, a cough. 13 */ 14 class Segmenter( 15 sampleRate: Int = 16_000, 16 blockMs: Int = 32, 17 private val floor: Float = 0.004f, 18 private val ratio: Float = 3f, 19 pauseMs: Int = 650, 20 minSpeechMs: Int = 450, 21 maxMs: Int = 11_000, 22 prerollMs: Int = 320, 23 ) { 24 val blockSamples = sampleRate * blockMs / 1000 25 private val pauseBlocks = (pauseMs + blockMs - 1) / blockMs 26 private val minSpeechBlocks = (minSpeechMs + blockMs - 1) / blockMs 27 private val maxSamples = sampleRate * maxMs / 1000 28 private val prerollBlocks = (prerollMs + blockMs - 1) / blockMs 29 30 private val partial = FloatArray(blockSamples) 31 private var partialCount = 0 32 private val preroll = ArrayDeque<FloatArray>() 33 private var open: MutableList<FloatArray>? = null 34 private var openSamples = 0 35 private var speechBlocks = 0 36 private var silenceBlocks = 0 37 private var noise = floor 38 39 /** Feeds the first [count] samples of [samples]. The pieces that closed go to [out]. */ 40 fun feed(samples: FloatArray, count: Int, out: MutableList<FloatArray>) { 41 var i = 0 42 while (i < count) { 43 val take = minOf(blockSamples - partialCount, count - i) 44 System.arraycopy(samples, i, partial, partialCount, take) 45 partialCount += take 46 i += take 47 if (partialCount == blockSamples) { 48 block(partial.copyOf(), out) 49 partialCount = 0 50 } 51 } 52 } 53 54 /** The open piece at the end of the sound, when it has enough speech; else null. */ 55 fun flush(): FloatArray? { 56 val blocks = open ?: return null 57 open = null 58 return if (speechBlocks >= minSpeechBlocks) join(blocks) else null 59 } 60 61 private fun block(block: FloatArray, out: MutableList<FloatArray>) { 62 val level = rms(block) 63 val speech = level > maxOf(floor, noise * ratio) 64 // The noise floor follows the quiet blocks: down at once, up slowly. 65 if (!speech) noise += (level - noise) * (if (level < noise) 0.5f else 0.05f) 66 67 val blocks = open 68 if (blocks == null) { 69 preroll.addLast(block) 70 while (preroll.size > prerollBlocks) preroll.removeFirst() 71 if (speech) { 72 open = preroll.toMutableList() 73 openSamples = preroll.sumOf { it.size } 74 preroll.clear() 75 speechBlocks = 1 76 silenceBlocks = 0 77 } 78 return 79 } 80 81 blocks.add(block) 82 openSamples += block.size 83 if (speech) { 84 speechBlocks++ 85 silenceBlocks = 0 86 } else { 87 silenceBlocks++ 88 } 89 when { 90 silenceBlocks >= pauseBlocks -> close(out, goesOn = false) 91 openSamples >= maxSamples -> close(out, goesOn = speech) 92 } 93 } 94 95 /** Closes the open piece. When the speech [goesOn], the next piece starts at once. */ 96 private fun close(out: MutableList<FloatArray>, goesOn: Boolean) { 97 val blocks = open ?: return 98 if (speechBlocks >= minSpeechBlocks) out.add(join(blocks)) 99 open = if (goesOn) mutableListOf() else null 100 openSamples = 0 101 speechBlocks = 0 102 silenceBlocks = 0 103 } 104 105 private fun join(blocks: List<FloatArray>): FloatArray { 106 val whole = FloatArray(blocks.sumOf { it.size }) 107 var at = 0 108 for (b in blocks) { 109 System.arraycopy(b, 0, whole, at, b.size) 110 at += b.size 111 } 112 return whole 113 } 114 115 private fun rms(block: FloatArray): Float { 116 var sum = 0.0 117 for (x in block) sum += x.toDouble() * x 118 return sqrt(sum / block.size).toFloat() 119 } 120 }