ChazzyG/Retrieval-based-Voice-Conversion-WebUI

1import os, traceback, sys, parselmouth2import librosa3import pyworld4from scipy.io import wavfile5import numpy as np, logging6 7logging.getLogger("numba").setLevel(logging.WARNING)8from multiprocessing import Process9 10exp_dir = sys.argv[1]11f = open("%s/extract_f0_feature.log" % exp_dir, "a+")12 13 14def printt(strr):15    print(strr)16    f.write("%s\n" % strr)17    f.flush()18 19 20n_p = int(sys.argv[2])21f0method = sys.argv[3]22 23 24class FeatureInput(object):25    def __init__(self, samplerate=16000, hop_size=160):26        self.fs = samplerate27        self.hop = hop_size28 29        self.f0_bin = 25630        self.f0_max = 1100.031        self.f0_min = 50.032        self.f0_mel_min = 1127 * np.log(1 + self.f0_min / 700)33        self.f0_mel_max = 1127 * np.log(1 + self.f0_max / 700)34 35    def compute_f0(self, path, f0_method):36        # default resample type of librosa.resample is "soxr_hq".37        # Quality: soxr_vhq > soxr_hq38        x, sr = librosa.load(path, self.fs)  # , res_type='soxr_vhq'39        p_len = x.shape[0] // self.hop40        assert sr == self.fs41        if f0_method == "pm":42            time_step = 160 / 16000 * 100043            f0_min = 5044            f0_max = 110045            f0 = (46                parselmouth.Sound(x, sr)47                .to_pitch_ac(48                    time_step=time_step / 1000,49                    voicing_threshold=0.6,50                    pitch_floor=f0_min,51                    pitch_ceiling=f0_max,52                )53                .selected_array["frequency"]54            )55            pad_size = (p_len - len(f0) + 1) // 256            if pad_size > 0 or p_len - len(f0) - pad_size > 0:57                f0 = np.pad(58                    f0, [[pad_size, p_len - len(f0) - pad_size]], mode="constant"59                )60        elif f0_method == "harvest":61            f0, t = pyworld.harvest(62                x.astype(np.double),63                fs=sr,64                f0_ceil=self.f0_max,65                f0_floor=self.f0_min,66                frame_period=1000 * self.hop / sr,67            )68            f0 = pyworld.stonemask(x.astype(np.double), f0, t, self.fs)69        elif f0_method == "dio":70            f0, t = pyworld.dio(71                x.astype(np.double),72                fs=sr,73                f0_ceil=self.f0_max,74                f0_floor=self.f0_min,75                frame_period=1000 * self.hop / sr,76            )77            f0 = pyworld.stonemask(x.astype(np.double), f0, t, self.fs)78        return f079 80    def coarse_f0(self, f0):81        f0_mel = 1127 * np.log(1 + f0 / 700)82        f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - self.f0_mel_min) * (83            self.f0_bin - 284        ) / (self.f0_mel_max - self.f0_mel_min) + 185 86        # use 0 or 187        f0_mel[f0_mel <= 1] = 188        f0_mel[f0_mel > self.f0_bin - 1] = self.f0_bin - 189        f0_coarse = np.rint(f0_mel).astype(np.int)90        assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (91            f0_coarse.max(),92            f0_coarse.min(),93        )94        return f0_coarse95 96    def go(self, paths, f0_method):97        if len(paths) == 0:98            printt("no-f0-todo")99        else:100            printt("todo-f0-%s" % len(paths))101            n = max(len(paths) // 5, 1)  # 每个进程最多打印5条102            for idx, (inp_path, opt_path1, opt_path2) in enumerate(paths):103                try:104                    if idx % n == 0:105                        printt("f0ing,now-%s,all-%s,-%s" % (idx, len(paths), inp_path))106                    if (107                        os.path.exists(opt_path1 + ".npy") == True108                        and os.path.exists(opt_path2 + ".npy") == True109                    ):110                        continue111                    featur_pit = self.compute_f0(inp_path, f0_method)112                    np.save(113                        opt_path2,114                        featur_pit,115                        allow_pickle=False,116                    )  # nsf117                    coarse_pit = self.coarse_f0(featur_pit)118                    np.save(119                        opt_path1,120                        coarse_pit,121                        allow_pickle=False,122                    )  # ori123                except:124                    printt("f0fail-%s-%s-%s" % (idx, inp_path, traceback.format_exc()))125 126 127if __name__ == "__main__":128    # exp_dir=r"E:\codes\py39\dataset\mi-test"129    # n_p=16130    # f = open("%s/log_extract_f0.log"%exp_dir, "w")131    printt(sys.argv)132    featureInput = FeatureInput()133    paths = []134    inp_root = "%s/1_16k_wavs" % (exp_dir)135    opt_root1 = "%s/2a_f0" % (exp_dir)136    opt_root2 = "%s/2b-f0nsf" % (exp_dir)137 138    os.makedirs(opt_root1, exist_ok=True)139    os.makedirs(opt_root2, exist_ok=True)140    for name in sorted(list(os.listdir(inp_root))):141        inp_path = "%s/%s" % (inp_root, name)142        if "spec" in inp_path:143            continue144        opt_path1 = "%s/%s" % (opt_root1, name)145        opt_path2 = "%s/%s" % (opt_root2, name)146        paths.append([inp_path, opt_path1, opt_path2])147 148    ps = []149    for i in range(n_p):150        p = Process(151            target=featureInput.go,152            args=(153                paths[i::n_p],154                f0method,155            ),156        )157        p.start()158        ps.append(p)159    for p in ps:160        p.join()161