From 867855ace9bf4b824cdfcb99251663a581573e7e Mon Sep 17 00:00:00 2001 From: litagin02 Date: Fri, 15 Mar 2024 11:10:23 +0900 Subject: [PATCH] Refactor, improve: use pathlib, allow audios other than wavs in resample --- resample.py | 92 ++++++++++++++++++++++++-------------------------- webui/train.py | 12 +++---- 2 files changed, 50 insertions(+), 54 deletions(-) diff --git a/resample.py b/resample.py index 5c3cc79..f9f1e3b 100644 --- a/resample.py +++ b/resample.py @@ -1,18 +1,19 @@ import argparse -import os -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor, as_completed from multiprocessing import cpu_count +from pathlib import Path +from typing import Any import librosa import pyloudnorm as pyln import soundfile +from numpy.typing import NDArray from tqdm import tqdm from config import config from style_bert_vits2.logging import logger from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT - DEFAULT_BLOCK_SIZE: float = 0.400 # seconds @@ -20,32 +21,37 @@ class BlockSizeException(Exception): pass -def normalize_audio(data, sr): +def normalize_audio(data: NDArray[Any], sr: int): meter = pyln.Meter(sr, block_size=DEFAULT_BLOCK_SIZE) # create BS.1770 meter try: loudness = meter.integrated_loudness(data) except ValueError as e: raise BlockSizeException(e) - # logger.info(f"loudness: {loudness}") + data = pyln.normalize.loudness(data, loudness, -23.0) return data -def process(item): - spkdir, wav_name, args = item - wav_path = os.path.join(args.in_dir, spkdir, wav_name) - if os.path.exists(wav_path) and wav_path.lower().endswith(".wav"): - wav, sr = librosa.load(wav_path, sr=args.sr) - if args.normalize: +def resample(file: Path, output_dir: Path, target_sr: int, normalize: bool, trim: bool): + try: + # librosaが読めるファイルかチェック + # wav以外にもmp3やoggやflacなども読める + wav: NDArray[Any] + sr: int + wav, sr = librosa.load(file, sr=target_sr) + if normalize: try: wav = normalize_audio(wav, sr) except BlockSizeException: + print("") logger.info( - f"Skip normalize due to less than {DEFAULT_BLOCK_SIZE} second audio: {wav_path}" + f"Skip normalize due to less than {DEFAULT_BLOCK_SIZE} second audio: {file}" ) - if args.trim: + if trim: wav, _ = librosa.effects.trim(wav, top_db=30) - soundfile.write(os.path.join(args.out_dir, spkdir, wav_name), wav, sr) + soundfile.write(output_dir / file.with_suffix(".wav").name, wav, sr) + except Exception as e: + logger.warning(f"Cannot load file, so skipping: {file}, {e}") if __name__ == "__main__": @@ -57,14 +63,14 @@ if __name__ == "__main__": help="sampling rate", ) parser.add_argument( - "--in_dir", + "--input_dir", "-i", type=str, default=config.resample_config.in_dir, help="path to source dir", ) parser.add_argument( - "--out_dir", + "--output_dir", "-o", type=str, default=config.resample_config.out_dir, @@ -88,46 +94,36 @@ if __name__ == "__main__": default=False, help="trim silence (start and end only)", ) - args, _ = parser.parse_known_args() - # autodl 无卡模式会识别出46个cpu + args = parser.parse_args() + if args.num_processes == 0: processes = cpu_count() - 2 if cpu_count() > 4 else 1 else: - processes = args.num_processes + processes: int = args.num_processes - tasks = [] + input_dir = Path(args.input_dir) + output_dir = Path(args.output_dir) + sr = int(args.sr) + normalize: bool = args.normalize + trim: bool = args.trim - for dirpath, _, filenames in os.walk(args.in_dir): - # 子级目录 - spk_dir = os.path.relpath(dirpath, args.in_dir) - spk_dir_out = os.path.join(args.out_dir, spk_dir) - if not os.path.isdir(spk_dir_out): - os.makedirs(spk_dir_out, exist_ok=True) - for filename in filenames: - if filename.lower().endswith(".wav"): - twople = (spk_dir, filename, args) - tasks.append(twople) + # 後でlibrosaに読ませて有効な音声ファイルかチェックするので、全てのファイルを取得 + original_files = [f for f in input_dir.rglob("*") if f.is_file()] - if len(tasks) == 0: - logger.error(f"No wav files found in {args.in_dir}") - raise ValueError(f"No wav files found in {args.in_dir}") + if len(original_files) == 0: + logger.error(f"No files found in {input_dir}") + raise ValueError(f"No files found in {input_dir}") - # pool = Pool(processes=processes) - # for _ in tqdm( - # pool.imap_unordered(process, tasks), file=SAFE_STDOUT, total=len(tasks) - # ): - # pass - - # pool.close() - # pool.join() + output_dir.mkdir(parents=True, exist_ok=True) with ThreadPoolExecutor(max_workers=processes) as executor: - _ = list( - tqdm( - executor.map(process, tasks), - total=len(tasks), - file=SAFE_STDOUT, - ) - ) + futures = [ + executor.submit(resample, file, output_dir, sr, normalize, trim) + for file in original_files + ] + for future in tqdm( + as_completed(futures), total=len(original_files), file=SAFE_STDOUT + ): + pass logger.info("Resampling Done!") diff --git a/webui/train.py b/webui/train.py index 6d72061..f7c1659 100644 --- a/webui/train.py +++ b/webui/train.py @@ -129,14 +129,14 @@ def initialize( def resample(model_name, normalize, trim, num_processes): logger.info("Step 2: start resampling...") dataset_path, _, _, _, _ = get_path(model_name) - in_dir = os.path.join(dataset_path, "raw") - out_dir = os.path.join(dataset_path, "wavs") + input_dir = os.path.join(dataset_path, "raw") + output_dir = os.path.join(dataset_path, "wavs") cmd = [ "resample.py", - "--in_dir", - in_dir, - "--out_dir", - out_dir, + "-i", + input_dir, + "-o", + output_dir, "--num_processes", str(num_processes), "--sr",