Spaces:
Running
Running
| # Copyright (c) 2023 Amphion. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import os | |
| import json | |
| import pickle | |
| import glob | |
| from collections import defaultdict | |
| from tqdm import tqdm | |
| from preprocessors import get_golden_samples_indexes | |
| TRAIN_MAX_NUM_EVERY_PERSON = 250 | |
| TEST_MAX_NUM_EVERY_PERSON = 25 | |
| def select_sample_idxs(): | |
| # =========== Train =========== | |
| with open(os.path.join(vctk_dir, "train.json"), "r") as f: | |
| raw_train = json.load(f) | |
| train_idxs = [] | |
| train_nums = defaultdict(int) | |
| for utt in tqdm(raw_train): | |
| idx = utt["index"] | |
| singer = utt["Singer"] | |
| if train_nums[singer] < TRAIN_MAX_NUM_EVERY_PERSON: | |
| train_idxs.append(idx) | |
| train_nums[singer] += 1 | |
| # =========== Test =========== | |
| with open(os.path.join(vctk_dir, "test.json"), "r") as f: | |
| raw_test = json.load(f) | |
| # golden test | |
| test_idxs = get_golden_samples_indexes( | |
| dataset_name="vctk", split="test", dataset_dir=vctk_dir | |
| ) | |
| test_nums = defaultdict(int) | |
| for idx in test_idxs: | |
| singer = raw_test[idx]["Singer"] | |
| test_nums[singer] += 1 | |
| for utt in tqdm(raw_test): | |
| idx = utt["index"] | |
| singer = utt["Singer"] | |
| if test_nums[singer] < TEST_MAX_NUM_EVERY_PERSON: | |
| test_idxs.append(idx) | |
| test_nums[singer] += 1 | |
| train_idxs.sort() | |
| test_idxs.sort() | |
| return train_idxs, test_idxs, raw_train, raw_test | |
| if __name__ == "__main__": | |
| root_path = "" | |
| vctk_dir = os.path.join(root_path, "vctk") | |
| sample_dir = os.path.join(root_path, "vctksample") | |
| os.makedirs(sample_dir, exist_ok=True) | |
| train_idxs, test_idxs, raw_train, raw_test = select_sample_idxs() | |
| print("#Train = {}, #Test = {}".format(len(train_idxs), len(test_idxs))) | |
| for split, chosen_idxs, utterances in zip( | |
| ["train", "test"], [train_idxs, test_idxs], [raw_train, raw_test] | |
| ): | |
| print( | |
| "#{} = {}, #chosen idx = {}\n".format( | |
| split, len(utterances), len(chosen_idxs) | |
| ) | |
| ) | |
| # Select features | |
| feat_files = glob.glob( | |
| "**/{}.pkl".format(split), root_dir=vctk_dir, recursive=True | |
| ) | |
| for file in tqdm(feat_files): | |
| raw_file = os.path.join(vctk_dir, file) | |
| new_file = os.path.join(sample_dir, file) | |
| new_dir = "/".join(new_file.split("/")[:-1]) | |
| os.makedirs(new_dir, exist_ok=True) | |
| if "mel_min" in file or "mel_max" in file: | |
| os.system("cp {} {}".format(raw_file, new_file)) | |
| continue | |
| with open(raw_file, "rb") as f: | |
| raw_feats = pickle.load(f) | |
| print("file: {}, #raw_feats = {}".format(file, len(raw_feats))) | |
| new_feats = [raw_feats[idx] for idx in chosen_idxs] | |
| with open(new_file, "wb") as f: | |
| pickle.dump(new_feats, f) | |
| # Utterance re-index | |
| news_utts = [utterances[idx] for idx in chosen_idxs] | |
| for i, utt in enumerate(news_utts): | |
| utt["Dataset"] = "vctksample" | |
| utt["index"] = i | |
| with open(os.path.join(sample_dir, "{}.json".format(split)), "w") as f: | |
| json.dump(news_utts, f, indent=4) | |