subrep-android/app/src/main/kotlin/com/honjimaku/subrep/Engine.kt

app/src/main/kotlin/com/honjimaku/subrep/Engine.kt (4605 bytes)
1 package com.honjimaku.subrep
2 
3 import java.util.concurrent.LinkedBlockingDeque
4 import kotlin.concurrent.thread
5 
6 /**
7  * Speech recognition: the sound goes through the [Segmenter], and each piece goes to the
8  * [Recognizer] (Whisper on the phone, or the cloud) on one worker thread. A piece that waits
9  * while the recognizer is behind by more than [QUEUE_LIMIT] pieces is dropped, the oldest first:
10  * a late caption helps nobody.
11  */
12 class Engine(
13     private val open: () -> Recognizer,
14     private val language: String,
15     private val onCaption: (Caption) -> Unit,
16     private val onLanguage: (String) -> Unit,
17     private val onState: (String) -> Unit,
18 ) {
19     private val segmenter = Segmenter()
20     private val queue = LinkedBlockingDeque<FloatArray>()
21     private val pieces = ArrayList<FloatArray>()
22     private var floats = FloatArray(0)
23     private var worker: Thread? = null
24 
25     @Volatile private var stopping = false
26 
27     /** How many pieces were dropped because Whisper was behind. */
28     @Volatile var late = 0
29         private set
30 
31     /** The time of the last piece divided by its length: below 1 keeps up. */
32     @Volatile var speed = 0f
33         private set
34 
35     fun start() {
36         onState(STATE_LOADING)
37         worker = thread(name = "whisper") { work() }
38     }
39 
40     /** The first [count] bytes of [frame], 16-bit PCM, from the capture thread. */
41     fun feed(frame: ByteArray, count: Int) {
42         if (floats.size < count / 2) floats = FloatArray(count / 2)
43         val n = Pcm.toFloat(frame, count, floats)
44         pieces.clear()
45         segmenter.feed(floats, n, pieces)
46         for (piece in pieces) offer(piece)
47     }
48 
49     /** Stops after the piece that Whisper works on, and the piece that was open. */
50     fun stop() {
51         stopping = true
52         segmenter.flush()?.let { queue.offerLast(it) }
53         queue.offerLast(END)
54         worker?.join(STOP_WAIT_MS)
55         worker?.interrupt()
56         worker = null
57     }
58 
59     private fun offer(piece: FloatArray) {
60         while (queue.size >= QUEUE_LIMIT) {
61             queue.pollFirst() ?: break
62             late++
63         }
64         queue.offerLast(piece)
65     }
66 
67     private fun work() {
68         val whisper = try {
69             open()
70         } catch (e: Exception) {
71             onState("${STATE_ERROR}${e.message}")
72             return
73         } catch (e: UnsatisfiedLinkError) {
74             onState("${STATE_ERROR}${e.message}")
75             return
76         }
77         onState(STATE_READY)
78         var lastLanguage = ""
79         var failed = false
80         try {
81             while (!Thread.currentThread().isInterrupted) {
82                 val piece = queue.takeFirst()
83                 if (piece === END) break
84                 val started = System.nanoTime()
85                 val text = try {
86                     Hallucinations.collapse(whisper.transcribe(piece, language)).trim()
87                 } catch (e: WhisperUnavailable) {
88                     failed = true
89                     onState("${STATE_ERROR}${e.message}")
90                     continue
91                 } catch (e: Cloud.PieceFailed) {
92                     // One lost piece, for example a short drop of the network. The next can work.
93                     failed = true
94                     onState("${STATE_ERROR}${e.message}")
95                     continue
96                 } catch (e: Cloud.OutOfHours) {
97                     // Each next piece would fail the same way.
98                     onState("${STATE_ERROR}${e.message}")
99                     break
100                 }
101                 if (failed) {
102                     failed = false
103                     onState(STATE_READY)
104                 }
105                 speed = (System.nanoTime() - started) / 1e9f / (piece.size / 16_000f)
106                 if (language == "auto") {
107                     val found = whisper.detectedLanguage
108                     if (found.isNotEmpty() && found != lastLanguage) {
109                         lastLanguage = found
110                         onLanguage(found)
111                     }
112                 }
113                 if (text.isEmpty() || Hallucinations.isHallucination(text)) continue
114                 onCaption(Caption("final", text))
115             }
116         } catch (_: InterruptedException) {
117             // Stopped while a piece waited.
118         } finally {
119             whisper.close()
120         }
121     }
122 
123     companion object {
124         const val STATE_LOADING = "loading"
125         const val STATE_READY = "ready"
126         const val STATE_ERROR = "error: "
127         private const val QUEUE_LIMIT = 3
128         private const val STOP_WAIT_MS = 8_000L
129         private val END = FloatArray(0)
130     }
131 }