File size: 3,695 Bytes
37f86d1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
124
125
126
127
128
129
130
131
132
import argparse
import os
from concurrent.futures import ThreadPoolExecutor

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()

    with ThreadPoolExecutor(max_workers=processes) as executor:
        _ = list(
            tqdm(
                executor.map(process, tasks),
                total=len(tasks),
                file=SAFE_STDOUT,
            )
        )

    logger.info("Resampling Done!")