mosi77/5

1from multiprocessing import cpu_count
2import os
3import sys
4
5from scipy import signal
6from scipy.io import wavfile
7import librosa
8import numpy as np
9
10now_directory = os.getcwd()
11sys.path.append(now_directory)
12
13from rvc.lib.utils import load_audio
14from rvc.train.slicer import Slicer
15
16experiment_directory = sys.argv[1]
17input_root = sys.argv[2]
18sampling_rate = int(sys.argv[3])
19percentage = float(sys.argv[4])
20num_processes = cpu_count()
21
22import multiprocessing
23
24
25class PreProcess:
26    def __init__(self, sr, exp_dir, per=3.0):
27        self.slicer = Slicer(
28            sr=sr,
29            threshold=-42,
30            min_length=1500,
31            min_interval=400,
32            hop_size=15,
33            max_sil_kept=500,
34        )
35        self.sr = sr
36        self.b_high, self.a_high = signal.butter(N=5, Wn=48, btype="high", fs=self.sr)
37        self.per = per
38        self.overlap = 0.3
39        self.tail = self.per + self.overlap
40        self.max_amplitude = 0.9
41        self.alpha = 0.75
42        self.exp_dir = exp_dir
43        self.gt_wavs_dir = f"{exp_dir}/0_gt_wavs"
44        self.wavs16k_dir = f"{exp_dir}/1_16k_wavs"
45        os.makedirs(self.exp_dir, exist_ok=True)
46        os.makedirs(self.gt_wavs_dir, exist_ok=True)
47        os.makedirs(self.wavs16k_dir, exist_ok=True)
48
49    def normalize_and_write(self, tmp_audio, idx0, idx1):
50        tmp_max = np.abs(tmp_audio).max()
51        if tmp_max > 2.5:
52            print(f"{idx0}-{idx1}-{tmp_max}-filtered")
53            return
54        tmp_audio = (tmp_audio / tmp_max * (self.max_amplitude * self.alpha)) + (
55            1 - self.alpha
56        ) * tmp_audio
57        wavfile.write(
58            f"{self.gt_wavs_dir}/{idx0}_{idx1}.wav",
59            self.sr,
60            tmp_audio.astype(np.float32),
61        )
62        tmp_audio = librosa.resample(
63            tmp_audio, orig_sr=self.sr, target_sr=16000
64        )  # , res_type="soxr_vhq"
65        wavfile.write(
66            f"{self.wavs16k_dir}/{idx0}_{idx1}.wav",
67            16000,
68            tmp_audio.astype(np.float32),
69        )
70
71    def process_audio(self, path, idx0):
72        try:
73            audio = load_audio(path, self.sr)
74            audio = signal.lfilter(self.b_high, self.a_high, audio)
75
76            idx1 = 0
77            for audio_segment in self.slicer.slice(audio):
78                i = 0
79                while 1:
80                    start = int(self.sr * (self.per - self.overlap) * i)
81                    i += 1
82                    if len(audio_segment[start:]) > self.tail * self.sr:
83                        tmp_audio = audio_segment[
84                            start : start + int(self.per * self.sr)
85                        ]
86                        self.normalize_and_write(tmp_audio, idx0, idx1)
87                        idx1 += 1
88                    else:
89                        tmp_audio = audio_segment[start:]
90                        idx1 += 1
91                        break
92                self.normalize_and_write(tmp_audio, idx0, idx1)
93        except Exception as error:
94            print(f"{path}: {error}")
95
96    def process_audio_multiprocessing(self, infos):
97        for path, idx0 in infos:
98            self.process_audio(path, idx0)
99
100    def process_audio_multiprocessing_input_directory(self, input_root, num_processes):
101        try:
102            infos = [
103                (f"{input_root}/{name}", idx)
104                for idx, name in enumerate(sorted(list(os.listdir(input_root))))
105            ]
106            processes = []
107            for i in range(num_processes):
108                p = multiprocessing.Process(
109                    target=self.process_audio_multiprocessing,
110                    args=(infos[i::num_processes],),
111                )
112                processes.append(p)
113                p.start()
114            for i in range(num_processes):
115                processes[i].join()
116        except Exception as error:
117            print(error)
118
119
120def preprocess_training_set(input_root, sr, num_processes, exp_dir, per):
121    pp = PreProcess(sr, exp_dir, per)
122    print("Starting preprocessing...")
123    pp.process_audio_multiprocessing_input_directory(input_root, num_processes)
124    print("Preprocessing completed!")
125
126
127if __name__ == "__main__":
128    preprocess_training_set(
129        input_root, sampling_rate, num_processes, experiment_directory, percentage
130    )
131