From 34f56ddb79556133830abcd8dfe37c1ead1d5a19 Mon Sep 17 00:00:00 2001 From: Matteo Fasulo Date: Thu, 9 Nov 2023 15:06:39 +0100 Subject: [PATCH] Add subtitle format and font color in word_options dict --- whisper/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/whisper/utils.py b/whisper/utils.py index 7a172c4..8f2df4e 100644 --- a/whisper/utils.py +++ b/whisper/utils.py @@ -115,12 +115,16 @@ class SubtitlesWriter(ResultWriter): max_line_count: Optional[int] = None, highlight_words: bool = False, max_words_per_line: Optional[int] = None, + subtitle_format: Optional[str] = None, + font_color: Optional[str] = None, ): options = options or {} max_line_width = max_line_width or options.get("max_line_width") max_line_count = max_line_count or options.get("max_line_count") highlight_words = highlight_words or options.get("highlight_words", False) max_words_per_line = max_words_per_line or options.get("max_words_per_line") + subtitle_format = subtitle_format or options.get("subtitle_format", None) + font_color = font_color or options.get("font_color", "#ffffff") preserve_segments = max_line_count is None or max_line_width is None max_line_width = max_line_width or 1000 max_words_per_line = max_words_per_line or 1000 @@ -195,7 +199,7 @@ class SubtitlesWriter(ResultWriter): yield start, end, "".join( [ - re.sub(r"^(\s*)(.*)$", r"\1\2", word) + re.sub(r"^(\s*)(.*)$", fr"\1<{subtitle_format}>\2", word) if j == i else word for j, word in enumerate(all_words)