From 3d36778e06a9be3a1a6891f7e28ee7c806c785d3 Mon Sep 17 00:00:00 2001
From: aka7774 <aka7774@users.noreply.github.com>
Date: Wed, 13 Mar 2024 16:07:09 +0900
Subject: [PATCH] Add time_suffix option

---
 slice.py         | 15 ++++++++++++++-
 webui_dataset.py |  9 ++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/slice.py b/slice.py
index 2d56427..382c32c 100644
--- a/slice.py
+++ b/slice.py
@@ -60,6 +60,7 @@ def split_wav(
     min_sec=2,
     max_sec=12,
     min_silence_dur_ms=700,
+    time_suffix=False,
 ):
     margin = 200  # ミリ秒単位で、音声の前後に余裕を持たせる
     speech_timestamps = get_stamps(
@@ -88,7 +89,11 @@ def split_wav(
         end_sample = int(end_ms / 1000 * sr)
         segment = data[start_sample:end_sample]
 
-        sf.write(os.path.join(target_dir, f"{file_name}-{i}.wav"), segment, sr)
+        if time_suffix:
+            file = f"{file_name}-{int(start_ms)}-{int(end_ms)}.wav"
+        else:
+            file = f"{file_name}-{i}.wav"
+        sf.write(os.path.join(target_dir, file), segment, sr)
         total_time_ms += end_ms - start_ms
         count += 1
 
@@ -123,6 +128,12 @@ if __name__ == "__main__":
         default=700,
         help="Silence above this duration (ms) is considered as a split point.",
     )
+    parser.add_argument(
+        "--time_suffix",
+        "-t",
+        action='store_true',
+        help="Make the filename end with -start_ms-end_ms when saving wav.",
+    )
     args = parser.parse_args()
 
     with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
@@ -134,6 +145,7 @@ if __name__ == "__main__":
     min_sec = args.min_sec
     max_sec = args.max_sec
     min_silence_dur_ms = args.min_silence_dur_ms
+    time_suffix = args.time_suffix
 
     wav_files = Path(input_dir).glob("**/*.wav")
     wav_files = list(wav_files)
@@ -151,6 +163,7 @@ if __name__ == "__main__":
             min_sec=min_sec,
             max_sec=max_sec,
             min_silence_dur_ms=min_silence_dur_ms,
+            time_suffix=time_suffix,
         )
         total_sec += time_sec
         total_count += count
diff --git a/webui_dataset.py b/webui_dataset.py
index fec7a9a..1bb1703 100644
--- a/webui_dataset.py
+++ b/webui_dataset.py
@@ -20,6 +20,7 @@ def do_slice(
     min_sec: float,
     max_sec: float,
     min_silence_dur_ms: int,
+    time_suffix: bool,
     input_dir: str,
 ):
     if model_name == "":
@@ -36,6 +37,8 @@ def do_slice(
         "--min_silence_dur_ms",
         str(min_silence_dur_ms),
     ]
+    if time_suffix:
+        cmd.append("--time_suffix")
     if input_dir != "":
         cmd += ["--input_dir", input_dir]
     # onnxの警告が出るので無視する
@@ -140,6 +143,10 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
                     step=100,
                     label="無音とみなして区切る最小の無音の長さ（ms）",
                 )
+                time_suffix = gr.Checkbox(
+                    value=False,
+                    label="WAVファイル名の末尾に元ファイルの時間範囲を付与する",
+                )
                 slice_button = gr.Button("スライスを実行")
             result1 = gr.Textbox(label="結果")
     with gr.Row():
@@ -174,7 +181,7 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
         result2 = gr.Textbox(label="結果")
     slice_button.click(
         do_slice,
-        inputs=[model_name, min_sec, max_sec, min_silence_dur_ms, input_dir],
+        inputs=[model_name, min_sec, max_sec, min_silence_dur_ms, time_suffix, input_dir],
         outputs=[result1],
     )
     transcribe_button.click(