1from __future__ import annotations
 2from typing import List
 3from dataclasses import dataclass
 4
 5
 6CONFIDENCE_THRESHOLD = 60
 7# predictions with lower confidence will be filtered out
 8
 9
10@dataclass
11class PredictedWord:
12    __slots__ = 'confidence', 'text'
13    confidence: int
14    text: str
15
16
17class PredictedFrame:
18    words: List[PredictedWord]
19
20    def __init__(self, pred_data: str):
21        self.words = []
22
23        block_current = 1
24        for line in pred_data.split('\n')[1:]:
25            tmp = line.split()
26            if len(tmp) < 12:
27                # no word is predicted
28                continue
29            _, _, block_num, *_, conf, text = tmp
30            block_num, conf = int(block_num), int(conf)
31
32            # handle line breaks
33            if block_current < block_num:
34                block_current = block_num
35                self.words.append(PredictedWord(0, '\n'))
36
37            if conf >= CONFIDENCE_THRESHOLD:
38                self.words.append(PredictedWord(conf, text))
39
40    @property
41    def confidence(self) -> int:
42        return sum(word.confidence for word in self.words)
43
44    @property
45    def text(self) -> str:
46        return ''.join(word.text + ' ' for word in self.words)
47