From ead77fab97e3abaec98a29e28f70874d619670e8 Mon Sep 17 00:00:00 2001
From: fatih <34196005+fcakyon@users.noreply.github.com>
Date: Mon, 26 Sep 2022 13:50:26 +0300
Subject: [PATCH] add srt subtitle export utility (#102)

* add srt subtitle export utility

* simplifying

Co-authored-by: Jong Wook Kim <jongwook@nyu.edu>
---
 whisper/transcribe.py |  6 +++++-
 whisper/utils.py      | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/whisper/transcribe.py b/whisper/transcribe.py
index 4c2123e36..82a705649 100644
--- a/whisper/transcribe.py
+++ b/whisper/transcribe.py
@@ -10,7 +10,7 @@
 from .audio import SAMPLE_RATE, N_FRAMES, HOP_LENGTH, pad_or_trim, log_mel_spectrogram
 from .decoding import DecodingOptions, DecodingResult
 from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE, get_tokenizer
-from .utils import exact_div, format_timestamp, optional_int, optional_float, str2bool, write_vtt
+from .utils import exact_div, format_timestamp, optional_int, optional_float, str2bool, write_vtt, write_srt
 
 if TYPE_CHECKING:
     from .model import Whisper
@@ -301,6 +301,10 @@ def cli():
         with open(os.path.join(output_dir, audio_basename + ".vtt"), "w", encoding="utf-8") as vtt:
             write_vtt(result["segments"], file=vtt)
 
+        # save SRT
+        with open(os.path.join(output_dir, audio_basename + ".srt"), "w", encoding="utf-8") as srt:
+            write_srt(result["segments"], file=srt)
+
 
 if __name__ == '__main__':
     cli()
diff --git a/whisper/utils.py b/whisper/utils.py
index af58834ee..c8f9c0dbe 100644
--- a/whisper/utils.py
+++ b/whisper/utils.py
@@ -27,7 +27,7 @@ def compression_ratio(text) -> float:
     return len(text) / len(zlib.compress(text.encode("utf-8")))
 
 
-def format_timestamp(seconds: float):
+def format_timestamp(seconds: float, always_include_hours: bool = False):
     assert seconds >= 0, "non-negative timestamp expected"
     milliseconds = round(seconds * 1000.0)
 
@@ -40,7 +40,8 @@ def format_timestamp(seconds: float):
     seconds = milliseconds // 1_000
     milliseconds -= seconds * 1_000
 
-    return (f"{hours}:" if hours > 0 else "") + f"{minutes:02d}:{seconds:02d}.{milliseconds:03d}"
+    hours_marker = f"{hours}:" if always_include_hours or hours > 0 else ""
+    return f"{hours_marker}{minutes:02d}:{seconds:02d}.{milliseconds:03d}"
 
 
 def write_vtt(transcript: Iterator[dict], file: TextIO):
@@ -52,3 +53,30 @@ def write_vtt(transcript: Iterator[dict], file: TextIO):
             file=file,
             flush=True,
         )
+
+
+def write_srt(transcript: Iterator[dict], file: TextIO):
+    """
+    Write a transcript to a file in SRT format.
+
+    Example usage:
+        from pathlib import Path
+        from whisper.utils import write_srt
+
+        result = transcribe(model, audio_path, temperature=temperature, **args)
+
+        # save SRT
+        audio_basename = Path(audio_path).stem
+        with open(Path(output_dir) / (audio_basename + ".srt"), "w", encoding="utf-8") as srt:
+            write_srt(result["segments"], file=srt)
+    """
+    for i, segment in enumerate(transcript, start=1):
+        # write srt lines
+        print(
+            f"{i}\n"
+            f"{format_timestamp(segment['start'], always_include_hours=True)} --> "
+            f"{format_timestamp(segment['end'], always_include_hours=True)}\n"
+            f"{segment['text'].strip().replace('-->', '->')}\n",
+            file=file,
+            flush=True,
+        )