Codecooker/rvcapi

1import sys, os, multiprocessing2from scipy import signal3 4now_dir = os.getcwd()5sys.path.append(now_dir)6 7inp_root = sys.argv[1]8sr = int(sys.argv[2])9n_p = int(sys.argv[3])10exp_dir = sys.argv[4]11noparallel = sys.argv[5] == "True"12import numpy as np, os, traceback13from slicer2 import Slicer14import librosa, traceback15from scipy.io import wavfile16import multiprocessing17from my_utils import load_audio18import tqdm19 20DoFormant = False21Quefrency = 1.022Timbre = 1.023 24mutex = multiprocessing.Lock()25f = open("%s/preprocess.log" % exp_dir, "a+")26 27 28def println(strr):29    mutex.acquire()30    print(strr)31    f.write("%s\n" % strr)32    f.flush()33    mutex.release()34 35 36class PreProcess:37    def __init__(self, sr, exp_dir):38        self.slicer = Slicer(39            sr=sr,40            threshold=-42,41            min_length=1500,42            min_interval=400,43            hop_size=15,44            max_sil_kept=500,45        )46        self.sr = sr47        self.bh, self.ah = signal.butter(N=5, Wn=48, btype="high", fs=self.sr)48        self.per = 3.049        self.overlap = 0.350        self.tail = self.per + self.overlap51        self.max = 0.952        self.alpha = 0.7553        self.exp_dir = exp_dir54        self.gt_wavs_dir = "%s/0_gt_wavs" % exp_dir55        self.wavs16k_dir = "%s/1_16k_wavs" % exp_dir56        os.makedirs(self.exp_dir, exist_ok=True)57        os.makedirs(self.gt_wavs_dir, exist_ok=True)58        os.makedirs(self.wavs16k_dir, exist_ok=True)59 60    def norm_write(self, tmp_audio, idx0, idx1):61        tmp_max = np.abs(tmp_audio).max()62        if tmp_max > 2.5:63            print("%s-%s-%s-filtered" % (idx0, idx1, tmp_max))64            return65        tmp_audio = (tmp_audio / tmp_max * (self.max * self.alpha)) + (66            1 - self.alpha67        ) * tmp_audio68        wavfile.write(69            "%s/%s_%s.wav" % (self.gt_wavs_dir, idx0, idx1),70            self.sr,71            tmp_audio.astype(np.float32),72        )73        tmp_audio = librosa.resample(74            tmp_audio, orig_sr=self.sr, target_sr=1600075        )  # , res_type="soxr_vhq"76        wavfile.write(77            "%s/%s_%s.wav" % (self.wavs16k_dir, idx0, idx1),78            16000,79            tmp_audio.astype(np.float32),80        )81 82    def pipeline(self, path, idx0):83        try:84            audio = load_audio(path, self.sr, DoFormant, Quefrency, Timbre)85            # zero phased digital filter cause pre-ringing noise...86            # audio = signal.filtfilt(self.bh, self.ah, audio)87            audio = signal.lfilter(self.bh, self.ah, audio)88 89            idx1 = 090            for audio in self.slicer.slice(audio):91                i = 092                while 1:93                    start = int(self.sr * (self.per - self.overlap) * i)94                    i += 195                    if len(audio[start:]) > self.tail * self.sr:96                        tmp_audio = audio[start : start + int(self.per * self.sr)]97                        self.norm_write(tmp_audio, idx0, idx1)98                        idx1 += 199                    else:100                        tmp_audio = audio[start:]101                        idx1 += 1102                        break103                self.norm_write(tmp_audio, idx0, idx1)104            # println("%s->Suc." % path)105        except:106            println("%s->%s" % (path, traceback.format_exc()))107 108    def pipeline_mp(self, infos, thread_n):109        for path, idx0 in tqdm.tqdm(110            infos, position=thread_n, leave=True, desc="thread:%s" % thread_n111        ):112            self.pipeline(path, idx0)113 114    def pipeline_mp_inp_dir(self, inp_root, n_p):115        try:116            infos = [117                ("%s/%s" % (inp_root, name), idx)118                for idx, name in enumerate(sorted(list(os.listdir(inp_root))))119            ]120            if noparallel:121                for i in range(n_p):122                    self.pipeline_mp(infos[i::n_p])123            else:124                ps = []125                for i in range(n_p):126                    p = multiprocessing.Process(127                        target=self.pipeline_mp, args=(infos[i::n_p], i)128                    )129                    ps.append(p)130                    p.start()131                for i in range(n_p):132                    ps[i].join()133        except:134            println("Fail. %s" % traceback.format_exc())135 136 137def preprocess_trainset(inp_root, sr, n_p, exp_dir):138    pp = PreProcess(sr, exp_dir)139    println("start preprocess")140    println(sys.argv)141    pp.pipeline_mp_inp_dir(inp_root, n_p)142    println("end preprocess")143 144 145if __name__ == "__main__":146    preprocess_trainset(inp_root, sr, n_p, exp_dir)147