add timestamps for words

example.py

@@ -1,73 +1,112 @@
-import sherpa_onnx
-import numpy as np
-import ujson
-from pydub import AudioSegment
-from datetime import datetime as dt
-
-
-recognizer = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(
-        model = "GigaAMv2_ctc_public.onnx",
-        tokens = "tokens.txt",
-        feature_dim=64,
-        num_threads=8,
-        sample_rate=8000,
-        decoding_method="greedy_search",
-        provider="CUDA",
-        debug = False,
-    )
-
-def get_np_array_samples_float32(audio_bytes: bytes, sample_width: int = 2) -> np.ndarray:
-    """
-    Преобразует аудио в байтах в массив float32.
-    :param audio_bytes: Аудиоданные в байтах.
-    :param sample_width: Размер одного сэмпла в байтах (обычно 2 для 16-битного аудио).
-    :return: Массив numpy с данными в формате float32.
-    """
-
-    # Определяем тип данных на основе sample_width
-    dtype = np.int16 if sample_width == 2 else np.int32
-
-    # Преобразуем байты в массив numpy
-    samples = np.frombuffer(audio_bytes, dtype=dtype)
-    samples_float32 = samples.astype(np.float32)
-    samples_float32 = samples_float32 / 32768
-
-    return samples_float32
-
-
-def simple_recognise(audio_data, ) -> dict:
-    """
-    Собираем токены в слова дополнительных вычислений не производит.
-    :param audio_data: Аудиоданные в формате Audiosegment (puDub).
-    """
-
-    stream = None
-    stream = recognizer.create_stream()
-    audio_data = audio_data.set_frame_rate(8000)
-    audio_data = audio_data[15000:]  # cut here to 15-18 secs
-    audio_data = audio_data.split_to_mono()[0]  # only mono allowed
-
-    # перевод в семплы для распознавания.
-    samples = get_np_array_samples_float32(audio_data.raw_data, audio_data.sample_width)
-
-    print(f'Audio length - {audio_data.duration_seconds} secs.')
-
-    # передали аудиофрагмент на распознавание
-    stream.accept_waveform(sample_rate=audio_data.frame_rate, waveform=samples)
-    recognizer.decode_stream(stream)
-
-    result = ujson.loads(str(stream.result))
-
-    return result
-
-
-if __name__ == '__main__':
-    file_path = "test.wav"
-    sound = AudioSegment.from_file(str(file_path))
-
-    time_start = dt.now()
-    asr_res = simple_recognise(audio_data=sound)
-    print(
-        f'Work time = {(dt.now()-time_start).total_seconds()}\n'
-    )
-    print(asr_res)
+import sherpa_onnx
+import numpy as np
+import ujson
+from pydub import AudioSegment
+from datetime import datetime as dt
+
+
+recognizer = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(
+        model = "GigaAMv2_ctc_public.onnx",
+        tokens = "tokens.txt",
+        feature_dim=64,
+        num_threads=8,
+        sample_rate=8000,
+        decoding_method="greedy_search",
+        provider="CUDA",
+        debug = False,
+    )
+
+def get_np_array_samples_float32(audio_bytes: bytes, sample_width: int = 2) -> np.ndarray:
+    """
+    Преобразует аудио в байтах в массив float32.
+    :param audio_bytes: Аудиоданные в байтах.
+    :param sample_width: Размер одного сэмпла в байтах (обычно 2 для 16-битного аудио).
+    :return: Массив numpy с данными в формате float32.
+    """
+
+    # Определяем тип данных на основе sample_width
+    dtype = np.int16 if sample_width == 2 else np.int32
+
+    # Преобразуем байты в массив numpy
+    samples = np.frombuffer(audio_bytes, dtype=dtype)
+    samples_float32 = samples.astype(np.float32)
+    samples_float32 = samples_float32 / 32768
+
+    return samples_float32
+
+
+def simple_recognise(audio_data, ) -> dict:
+    """
+    :param audio_data: Аудиоданные в формате Audiosegment (puDub).
+    """
+
+    stream = None
+    stream = recognizer.create_stream()
+    audio_data = audio_data.set_frame_rate(8000)
+    audio_data = audio_data[:]  # cut here to 15-18 secs
+    audio_data = audio_data.split_to_mono()[0]  # only mono allowed
+
+    # перевод в семплы для распознавания.
+    samples = get_np_array_samples_float32(audio_data.raw_data, audio_data.sample_width)
+
+    print(f'Audio length - {audio_data.duration_seconds} secs.')
+
+    # передали аудиофрагмент на распознавание
+    stream.accept_waveform(sample_rate=audio_data.frame_rate, waveform=samples)
+    recognizer.decode_stream(stream)
+
+    result = ujson.loads(str(stream.result))
+
+    return result
+
+def process_gigaam_asr(input_json, time_shift = 0.0):
+    """
+        Собираем токены в слова дополнительных вычислений не производим.
+        :param input_json: json - результат работы stream.result
+        :param input_json: time_shift - так как на вход логично будут приходить чанки,
+            то для каждого чанка передаём его начало от начала записи.
+    """
+    # Парсим JSON
+    data = input_json
+
+    # Формируем шаблон результата
+    result = {"data": {"result": [], "text": ""}}
+
+    # Собираем слова из токенов
+    words = []
+    current_word = ""
+    start_time, end_time = 0.0, 0.0
+
+    for i, token in enumerate(data['tokens']):
+        if token != ' ':
+            if current_word == "":
+                start_time = round((data['timestamps'][i]+time_shift), 3)
+            current_word += token
+            end_time = round((data['timestamps'][i]+time_shift), 3)
+        else:
+            if current_word != "":
+                words.append({'word': current_word, 'start': start_time, 'end': end_time})
+                current_word = ""
+
+    # Добавляем последнее слово, если оно есть
+    if current_word != "":
+        words.append({'word': current_word, 'start': start_time, 'end': end_time})
+
+    # Формируем итоговый массив
+    result['data'] = {
+        'result': [{'start': word['start'], 'end': word['end'], 'word': word['word']} for word in words],
+        'text': data['text']
+    }
+    return result
+
+if __name__ == '__main__':
+    file_path = "example.wav"
+    sound = AudioSegment.from_file(str(file_path))
+
+    time_start = dt.now()
+    asr_res = simple_recognise(audio_data=sound)
+    res_w_word_timestamp = process_gigaam_asr(asr_res)
+    print(
+        f'Work time = {(dt.now()-time_start).total_seconds()}\n'
+    )
+    print(res_w_word_timestamp)