Generate subtitles (#315)
This commit is contained in:
Fangjun Kuang
386
python-api-examples/generate-subtitles.py
Executable file
386
python-api-examples/generate-subtitles.py
Executable file
@@ -0,0 +1,386 @@
|
||||
#!/usr/bin/env python3
|
||||
#
|
||||
# Copyright (c) 2023 Xiaomi Corporation
|
||||
|
||||
"""
|
||||
This file demonstrates how to use sherpa-onnx Python APIs to generate
|
||||
subtitles.
|
||||
|
||||
Supported file formats are those supported by ffmpeg; for instance,
|
||||
*.mov, *.mp4, *.wav, etc.
|
||||
|
||||
Note that you need a non-streaming model for this script.
|
||||
|
||||
(1) For paraformer
|
||||
|
||||
./python-api-examples/generate-subtitles.py \
|
||||
--silero-vad-model=/path/to/silero_vad.onnx \
|
||||
--tokens=/path/to/tokens.txt \
|
||||
--paraformer=/path/to/paraformer.onnx \
|
||||
--num-threads=2 \
|
||||
--decoding-method=greedy_search \
|
||||
--debug=false \
|
||||
--sample-rate=16000 \
|
||||
--feature-dim=80 \
|
||||
/path/to/test.mp4
|
||||
|
||||
(2) For transducer models from icefall
|
||||
|
||||
./python-api-examples/generate-subtitles.py \
|
||||
--silero-vad-model=/path/to/silero_vad.onnx \
|
||||
--tokens=/path/to/tokens.txt \
|
||||
--encoder=/path/to/encoder.onnx \
|
||||
--decoder=/path/to/decoder.onnx \
|
||||
--joiner=/path/to/joiner.onnx \
|
||||
--num-threads=2 \
|
||||
--decoding-method=greedy_search \
|
||||
--debug=false \
|
||||
--sample-rate=16000 \
|
||||
--feature-dim=80 \
|
||||
/path/to/test.mp4
|
||||
|
||||
(3) For Whisper models
|
||||
|
||||
./python-api-examples/generate-subtitles.py \
|
||||
--silero-vad-model=/path/to/silero_vad.onnx \
|
||||
--whisper-encoder=./sherpa-onnx-whisper-base.en/base.en-encoder.int8.onnx \
|
||||
--whisper-decoder=./sherpa-onnx-whisper-base.en/base.en-decoder.int8.onnx \
|
||||
--tokens=./sherpa-onnx-whisper-base.en/base.en-tokens.txt \
|
||||
--whisper-task=transcribe \
|
||||
--num-threads=2 \
|
||||
/path/to/test.mp4
|
||||
|
||||
Please refer to
|
||||
https://k2-fsa.github.io/sherpa/onnx/index.html
|
||||
to install sherpa-onnx and to download non-streaming pre-trained models
|
||||
used in this file.
|
||||
"""
|
||||
import argparse
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import sherpa_onnx
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--silero-vad-model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to silero_vad.onnx",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--tokens",
|
||||
type=str,
|
||||
help="Path to tokens.txt",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--encoder",
|
||||
default="",
|
||||
type=str,
|
||||
help="Path to the transducer encoder model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--decoder",
|
||||
default="",
|
||||
type=str,
|
||||
help="Path to the transducer decoder model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--joiner",
|
||||
default="",
|
||||
type=str,
|
||||
help="Path to the transducer joiner model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--paraformer",
|
||||
default="",
|
||||
type=str,
|
||||
help="Path to the model.onnx from Paraformer",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--num-threads",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of threads for neural network computation",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--whisper-encoder",
|
||||
default="",
|
||||
type=str,
|
||||
help="Path to whisper encoder model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--whisper-decoder",
|
||||
default="",
|
||||
type=str,
|
||||
help="Path to whisper decoder model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--whisper-language",
|
||||
default="",
|
||||
type=str,
|
||||
help="""It specifies the spoken language in the input file.
|
||||
Example values: en, fr, de, zh, jp.
|
||||
Available languages for multilingual models can be found at
|
||||
https://github.com/openai/whisper/blob/main/whisper/tokenizer.py#L10
|
||||
If not specified, we infer the language from the input audio file.
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--whisper-task",
|
||||
default="transcribe",
|
||||
choices=["transcribe", "translate"],
|
||||
type=str,
|
||||
help="""For multilingual models, if you specify translate, the output
|
||||
will be in English.
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--decoding-method",
|
||||
type=str,
|
||||
default="greedy_search",
|
||||
help="""Valid values are greedy_search and modified_beam_search.
|
||||
modified_beam_search is valid only for transducer models.
|
||||
""",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
type=bool,
|
||||
default=False,
|
||||
help="True to show debug messages when loading modes.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--sample-rate",
|
||||
type=int,
|
||||
default=16000,
|
||||
help="""Sample rate of the feature extractor. Must match the one
|
||||
expected by the model. Note: The input sound files can have a
|
||||
different sample rate from this argument.""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--feature-dim",
|
||||
type=int,
|
||||
default=80,
|
||||
help="Feature dimension. Must match the one expected by the model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"sound_file",
|
||||
type=str,
|
||||
help="The input sound file to generate subtitles ",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def assert_file_exists(filename: str):
|
||||
assert Path(filename).is_file(), (
|
||||
f"{filename} does not exist!\n"
|
||||
"Please refer to "
|
||||
"https://k2-fsa.github.io/sherpa/onnx/pretrained_models/index.html to download it"
|
||||
)
|
||||
|
||||
|
||||
def create_recognizer(args) -> sherpa_onnx.OfflineRecognizer:
|
||||
if args.encoder:
|
||||
assert len(args.paraformer) == 0, args.paraformer
|
||||
assert len(args.whisper_encoder) == 0, args.whisper_encoder
|
||||
assert len(args.whisper_decoder) == 0, args.whisper_decoder
|
||||
|
||||
assert_file_exists(args.encoder)
|
||||
assert_file_exists(args.decoder)
|
||||
assert_file_exists(args.joiner)
|
||||
|
||||
recognizer = sherpa_onnx.OfflineRecognizer.from_transducer(
|
||||
encoder=args.encoder,
|
||||
decoder=args.decoder,
|
||||
joiner=args.joiner,
|
||||
tokens=args.tokens,
|
||||
num_threads=args.num_threads,
|
||||
sample_rate=args.sample_rate,
|
||||
feature_dim=args.feature_dim,
|
||||
decoding_method=args.decoding_method,
|
||||
debug=args.debug,
|
||||
)
|
||||
elif args.paraformer:
|
||||
assert len(args.whisper_encoder) == 0, args.whisper_encoder
|
||||
assert len(args.whisper_decoder) == 0, args.whisper_decoder
|
||||
|
||||
assert_file_exists(args.paraformer)
|
||||
|
||||
recognizer = sherpa_onnx.OfflineRecognizer.from_paraformer(
|
||||
paraformer=args.paraformer,
|
||||
tokens=args.tokens,
|
||||
num_threads=args.num_threads,
|
||||
sample_rate=args.sample_rate,
|
||||
feature_dim=args.feature_dim,
|
||||
decoding_method=args.decoding_method,
|
||||
debug=args.debug,
|
||||
)
|
||||
elif args.whisper_encoder:
|
||||
assert_file_exists(args.whisper_encoder)
|
||||
assert_file_exists(args.whisper_decoder)
|
||||
|
||||
recognizer = sherpa_onnx.OfflineRecognizer.from_whisper(
|
||||
encoder=args.whisper_encoder,
|
||||
decoder=args.whisper_decoder,
|
||||
tokens=args.tokens,
|
||||
num_threads=args.num_threads,
|
||||
decoding_method=args.decoding_method,
|
||||
debug=args.debug,
|
||||
language=args.whisper_language,
|
||||
task=args.whisper_task,
|
||||
)
|
||||
else:
|
||||
raise ValueError("Please specify at least one model")
|
||||
|
||||
return recognizer
|
||||
|
||||
|
||||
@dataclass
|
||||
class Segment:
|
||||
start: float
|
||||
duration: float
|
||||
text: str = ""
|
||||
|
||||
@property
|
||||
def end(self):
|
||||
return self.start + self.duration
|
||||
|
||||
def __str__(self):
|
||||
s = f"{timedelta(seconds=self.start)}"[:-3]
|
||||
s += " --> "
|
||||
s += f"{timedelta(seconds=self.end)}"[:-3]
|
||||
s = s.replace(".", ",")
|
||||
s += "\n"
|
||||
s += self.text
|
||||
return s
|
||||
|
||||
|
||||
def main():
|
||||
args = get_args()
|
||||
assert_file_exists(args.tokens)
|
||||
assert_file_exists(args.silero_vad_model)
|
||||
|
||||
assert args.num_threads > 0, args.num_threads
|
||||
|
||||
if not Path(args.sound_file).is_file():
|
||||
raise ValueError(f"{args.sound_file} does not exist")
|
||||
|
||||
assert (
|
||||
args.sample_rate == 16000
|
||||
), f"Only sample rate 16000 is supported.Given: {args.sample_rate}"
|
||||
|
||||
recognizer = create_recognizer(args)
|
||||
|
||||
ffmpeg_cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
args.sound_file,
|
||||
"-f",
|
||||
"s16le",
|
||||
"-acodec",
|
||||
"pcm_s16le",
|
||||
"-ac",
|
||||
"1",
|
||||
"-ar",
|
||||
str(args.sample_rate),
|
||||
"-",
|
||||
]
|
||||
|
||||
process = subprocess.Popen(
|
||||
ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL
|
||||
)
|
||||
|
||||
frames_per_read = int(args.sample_rate * 100) # 100 second
|
||||
|
||||
stream = recognizer.create_stream()
|
||||
|
||||
config = sherpa_onnx.VadModelConfig()
|
||||
config.silero_vad.model = args.silero_vad_model
|
||||
config.silero_vad.min_silence_duration = 0.25
|
||||
config.sample_rate = args.sample_rate
|
||||
|
||||
window_size = config.silero_vad.window_size
|
||||
|
||||
buffer = []
|
||||
vad = sherpa_onnx.VoiceActivityDetector(config, buffer_size_in_seconds=100)
|
||||
|
||||
segment_list = []
|
||||
|
||||
print("Started!")
|
||||
|
||||
# TODO(fangjun): Support multithreads
|
||||
while True:
|
||||
# *2 because int16_t has two bytes
|
||||
data = process.stdout.read(frames_per_read * 2)
|
||||
if not data:
|
||||
break
|
||||
|
||||
samples = np.frombuffer(data, dtype=np.int16)
|
||||
samples = samples.astype(np.float32) / 32768
|
||||
|
||||
buffer = np.concatenate([buffer, samples])
|
||||
while len(buffer) > window_size:
|
||||
vad.accept_waveform(buffer[:window_size])
|
||||
buffer = buffer[window_size:]
|
||||
|
||||
streams = []
|
||||
segments = []
|
||||
while not vad.empty():
|
||||
segment = Segment(
|
||||
start=vad.front.start / args.sample_rate,
|
||||
duration=len(vad.front.samples) / args.sample_rate,
|
||||
)
|
||||
segments.append(segment)
|
||||
|
||||
stream = recognizer.create_stream()
|
||||
stream.accept_waveform(args.sample_rate, vad.front.samples)
|
||||
|
||||
streams.append(stream)
|
||||
|
||||
vad.pop()
|
||||
|
||||
recognizer.decode_streams(streams)
|
||||
for seg, stream in zip(segments, streams):
|
||||
seg.text = stream.result.text
|
||||
segment_list.append(seg)
|
||||
|
||||
srt_filename = Path(args.sound_file).with_suffix(".srt")
|
||||
with open(srt_filename, "w", encoding="utf-8") as f:
|
||||
for i, seg in enumerate(segment_list):
|
||||
print(i + 1, file=f)
|
||||
print(seg, file=f)
|
||||
print("", file=f)
|
||||
|
||||
print(f"Saved to {srt_filename}")
|
||||
print("Done!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if shutil.which("ffmpeg") is None:
|
||||
sys.exit("Please install ffmpeg first!")
|
||||
main()
|
||||
Reference in New Issue
Block a user