File size: 3,436 Bytes
8dfbf56
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import argparse
import os
from multiprocessing import Pool, cpu_count

import librosa
import pyloudnorm as pyln
import soundfile
from tqdm import tqdm

from common.log import logger
from common.stdout_wrapper import SAFE_STDOUT
from config import config

DEFAULT_BLOCK_SIZE: float = 0.400  # seconds


class BlockSizeException(Exception):
    pass


def normalize_audio(data, sr):
    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:
            try:
                wav = normalize_audio(wav, sr)
            except BlockSizeException:
                logger.info(
                    f"Skip normalize due to less than {DEFAULT_BLOCK_SIZE} second audio: {wav_path}"
                )
        if args.trim:
            wav, _ = librosa.effects.trim(wav, top_db=30)
        soundfile.write(os.path.join(args.out_dir, spkdir, wav_name), wav, sr)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--sr",
        type=int,
        default=config.resample_config.sampling_rate,
        help="sampling rate",
    )
    parser.add_argument(
        "--in_dir",
        "-i",
        type=str,
        default=config.resample_config.in_dir,
        help="path to source dir",
    )
    parser.add_argument(
        "--out_dir",
        "-o",
        type=str,
        default=config.resample_config.out_dir,
        help="path to target dir",
    )
    parser.add_argument(
        "--num_processes",
        type=int,
        default=4,
        help="cpu_processes",
    )
    parser.add_argument(
        "--normalize",
        action="store_true",
        default=False,
        help="loudness normalize audio",
    )
    parser.add_argument(
        "--trim",
        action="store_true",
        default=False,
        help="trim silence (start and end only)",
    )
    args, _ = parser.parse_known_args()
    # autodl 无卡模式会识别出46个cpu
    if args.num_processes == 0:
        processes = cpu_count() - 2 if cpu_count() > 4 else 1
    else:
        processes = args.num_processes

    tasks = []

    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)

    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}")

    pool = Pool(processes=processes)
    for _ in tqdm(
        pool.imap_unordered(process, tasks), file=SAFE_STDOUT, total=len(tasks)
    ):
        pass

    pool.close()
    pool.join()

    logger.info("Resampling Done!")