Add args for Dataset and improve

slice.py

@@ -1,11 +1,13 @@
 import argparse
 import os
 import shutil
+from pathlib import Path
 
 import soundfile as sf
 import torch
 from tqdm import tqdm
 
+from common.log import logger
 from common.stdout_wrapper import SAFE_STDOUT
 from resample import normalize_audio
 
@@ -18,40 +20,53 @@ vad_model, utils = torch.hub.load(
 (get_speech_timestamps, _, read_audio, *_) = utils
 
 
-def get_stamps(audio_file, min_silence_dur_ms=700, min_sec=2):
+def get_stamps(
+    audio_file, min_silence_dur_ms: int = 700, min_sec: float = 2, max_sec: float = 12
+):
     """
-    min_silence_dur_ms:
+    min_silence_dur_ms: int (ミリ秒):
         このミリ秒数以上を無音だと判断する。
         逆に、この秒数以下の無音区間では区切られない。
         小さくすると、音声がぶつ切りに小さくなりすぎ、
         大きくすると音声一つ一つが長くなりすぎる。
         データセットによってたぶん要調整。
-    min_sec:
-        この秒数より小さい発話は無視する。TTSのためには2秒未満は切り捨てたほうがいいかも。
+    min_sec: float (秒):
+        この秒数より小さい発話は無視する。
+    max_sec: float (秒):
+        この秒数より大きい発話は無視する。
     """
 
     sampling_rate = 16000  # 16kHzか8kHzのみ対応
 
+    min_ms = int(min_sec * 1000)
+
     wav = read_audio(audio_file, sampling_rate=sampling_rate)
     speech_timestamps = get_speech_timestamps(
         wav,
         vad_model,
         sampling_rate=sampling_rate,
         min_silence_duration_ms=min_silence_dur_ms,
-        min_speech_duration_ms=min_sec * 1000,
+        min_speech_duration_ms=min_ms,
+        max_speech_duration_s=max_sec,
     )
 
     return speech_timestamps
 
 
 def split_wav(
-    audio_file, target_dir="raw", max_sec=12, min_silence_dur_ms=700, min_sec=2, normalize=False
+    audio_file,
+    target_dir="raw",
+    min_sec=2,
+    max_sec=12,
+    min_silence_dur_ms=700,
+    normalize=False,
 ):
     margin = 200  # ミリ秒単位で、音声の前後に余裕を持たせる
-    upper_bound_ms = max_sec * 1000  # これ以上の長さの音声は無視する
-
     speech_timestamps = get_stamps(
-        audio_file, min_silence_dur_ms=min_silence_dur_ms, min_sec=min_sec
+        audio_file,
+        min_silence_dur_ms=min_silence_dur_ms,
+        min_sec=min_sec,
+        max_sec=max_sec,
     )
 
     data, sr = sf.read(audio_file)
@@ -67,8 +82,6 @@ def split_wav(
     for i, ts in enumerate(speech_timestamps):
         start_ms = max(ts["start"] / 16 - margin, 0)
         end_ms = min(ts["end"] / 16 + margin, total_ms)
-        if end_ms - start_ms > upper_bound_ms:
-            continue
 
         start_sample = int(start_ms / 1000 * sr)
         end_sample = int(end_ms / 1000 * sr)
@@ -85,12 +98,36 @@ def split_wav(
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--max_sec", "-M", type=int, default=12)
-    parser.add_argument("--min_sec", "-m", type=int, default=2)
-    parser.add_argument("--min_silence_dur_ms", "-s", type=int, default=700)
-    parser.add_argument("--input_dir", "-i", type=str, default="inputs")
-    parser.add_argument("--output_dir", "-t", type=str, default="raw")
-    parser.add_argument("--normalize", action="store_true")
+    parser.add_argument(
+        "--min_sec", "-m", type=float, default=2, help="Minimum seconds of a slice"
+    )
+    parser.add_argument(
+        "--max_sec", "-M", type=float, default=12, help="Maximum seconds of a slice"
+    )
+    parser.add_argument(
+        "--input_dir",
+        "-i",
+        type=str,
+        default="inputs",
+        help="Directory of input wav files",
+    )
+    parser.add_argument(
+        "--output_dir",
+        "-t",
+        type=str,
+        default="raw",
+        help="Directory of output wav files",
+    )
+    parser.add_argument(
+        "--normalize", action="store_true", help="Whether to normalize loudness"
+    )
+    parser.add_argument(
+        "--min_silence_dur_ms",
+        "-s",
+        type=int,
+        default=700,
+        help="Silence above this duration (ms) is considered as a split point.",
+    )
     args = parser.parse_args()
 
     input_dir = args.input_dir
@@ -100,25 +137,22 @@ if __name__ == "__main__":
     min_silence_dur_ms = args.min_silence_dur_ms
     normalize = args.normalize
 
-    wav_files = [
-        os.path.join(input_dir, f)
-        for f in os.listdir(input_dir)
-        if f.lower().endswith(".wav")
-    ]
-    if os.path.exists(output_dir):  # ディレクトリを削除
-        print(f"{output_dir}フォルダが存在するので、削除します。")
+    wav_files = Path(input_dir).glob("**/*.wav")
+    wav_files = list(wav_files)
+    if os.path.exists(output_dir):
+        logger.warning(f"Output directory {output_dir} already exists, deleting...")
         shutil.rmtree(output_dir)
 
     total_sec = 0
     for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
         time_sec = split_wav(
-            wav_file,
-            output_dir,
-            max_sec=max_sec,
+            audio_file=str(wav_file),
+            target_dir=output_dir,
             min_sec=min_sec,
+            max_sec=max_sec,
             min_silence_dur_ms=min_silence_dur_ms,
-            normalize=normalize
+            normalize=normalize,
         )
         total_sec += time_sec
 
-    print(f"Done! Total time: {total_sec / 60:.2f} min.")
+    logger.info(f"Slice done! Total time: {total_sec / 60:.2f} min.")