aboutsummaryrefslogtreecommitdiffstats
path: root/scripts/02_transcribe_audio.py
blob: 3e513e4b41f2a2c656a50513e68f4efc6197f3b6 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#!/usr/local/bin/python3.12

from vosk import KaldiRecognizer, Model
import os.path
import json
import os


FRAME = 4_000

SAMPLE_RATE = 16_000

MODELS = list(
  map(
    lambda model: os.path.join(os.getcwd(), 'third-party/vosk-models', model),
    [
      'vosk-model-en-us-0.22-lgraph', # English (United States)
      'vosk-model-small-ru-0.22', # Russian
      'vosk-model-small-uk-v3-small', # Ukrainian
    ]
  )
)

DATASET = os.path.join(os.getcwd(), 'files/dataset.json')

AUDIOS = os.path.join(os.getcwd(), 'files/audios')


def score_transcription(result: dict) -> float:
  words = result.get("result", [])
  values = list(
    map(lambda word: word.get("conf", 0.0), words)
  )
  try:
    return sum(values) / len(values)
  except ZeroDivisionError:
    return 0.0


def main() -> None:
  recognizers = []
  transcriptions = []

  for model in MODELS:
    model = Model(model)
    recognizer = KaldiRecognizer(model, SAMPLE_RATE)
    recognizer.SetWords(True)
    recognizers.append(recognizer)

  for child in os.listdir(AUDIOS):
    if child == ".gitkeep": continue

    file = os.path.join(AUDIOS, child)
    if not os.path.isfile(file): continue

    with open(file, 'rb') as audio:
      while True:
        frame = audio.read(FRAME)
        if len(frame) == 0:
          break
        for recognizer in recognizers:
          recognizer.AcceptWaveform(frame)

      results = []
      for i, recognizer in enumerate(recognizers):
        result = json.loads(recognizer.FinalResult())
        if result.get("text", "") == "": continue

        score = score_transcription(result)
        if score == 0.0: continue

        results.append({'model': MODELS[i], 'score': score, 'result': result})

      if len(results) == 0:
        continue
      results = list(
        sorted(results, key=lambda result: result['score'], reverse=True)
      )

      transcription = results[0]
      transcription.update({'video': child})

      transcriptions.append(transcription)

  dataset = []
  for transcription in transcriptions:
    words = transcription['result']['result']
    message = []
    for i, word in enumerate(words):
      message.append(f'[{i + 1}]{word['word']}')
    dataset.append({'user': ' '.join(message), 'assistant': None})

  with open(DATASET, 'w') as file:
    json.dump(dataset, file, ensure_ascii=False, indent=2)


if __name__ == "__main__":
  main()