Spaces:
Running
Running
| from sklearn.cluster import * | |
| import os | |
| import numpy as np | |
| from config import config | |
| import yaml | |
| from multiprocessing import Pool | |
| from tqdm import tqdm | |
| def process_speaker(speaker): | |
| embs = [] | |
| wavnames = [] | |
| print("\nspeaker: " + speaker) | |
| for file in filelist_dict[speaker]: | |
| try: | |
| embs.append( | |
| np.expand_dims(np.load(f"{os.path.splitext(file)[0]}.emo.npy"), axis=0) | |
| ) | |
| wavnames.append(os.path.basename(file)) | |
| except Exception as e: | |
| print(e) | |
| x = np.concatenate(embs, axis=0) | |
| x = np.squeeze(x) | |
| # 聚类算法类的数量 | |
| n_clusters = config.emo_cluster_config.n_clusters | |
| method = config.emo_cluster_config.method | |
| if method == "b": | |
| model = Birch(n_clusters=n_clusters, threshold=0.2) | |
| elif method == "s": | |
| model = SpectralClustering(n_clusters=n_clusters) | |
| elif method == "a": | |
| model = AgglomerativeClustering(n_clusters=n_clusters) | |
| else: | |
| model = KMeans(n_clusters=n_clusters, random_state=42) | |
| # 可以自行尝试各种不同的聚类算法 | |
| y_predict = model.fit_predict(x) | |
| classes = [[] for i in range(y_predict.max() + 1)] | |
| for idx, wavname in enumerate(wavnames): | |
| classes[y_predict[idx]].append(wavname) | |
| yml_result = {} | |
| yml_result[speaker] = {} | |
| os.makedirs( | |
| os.path.join(config.dataset_path, f"emo_clustering/{speaker}"), exist_ok=True | |
| ) | |
| for i in range(y_predict.max() + 1): | |
| class_length = len(classes[i]) | |
| print("类别:", i, "本类中样本数量:", class_length) | |
| yml_result[speaker][f"class{i}"] = [] | |
| for j in range(config.emo_cluster_config.n_samples): | |
| if j >= class_length: | |
| break | |
| print(classes[i][j]) | |
| yml_result[speaker][f"class{i}"].append(classes[i][j]) | |
| if hasattr(model, "cluster_centers_") and config.emo_cluster_config.save_center: | |
| centers = model.cluster_centers_ | |
| filename = os.path.join( | |
| config.dataset_path, f"emo_clustering/{speaker}/cluster_center_{i}.npy" | |
| ) | |
| # 保存中心 | |
| np.save(filename, centers[i]) | |
| elif config.emo_cluster_config.save_center: | |
| labels = model.labels_ | |
| centers = np.array([X[labels == i].mean(0) for i in range(n_clusters)]) | |
| filename = os.path.join( | |
| config.dataset_path, f"emo_clustering/{speaker}/cluster_center_{i}.npy" | |
| ) | |
| np.save(filename, centers[i]) | |
| return yml_result | |
| if __name__ == "__main__": | |
| filelist_dict = {} | |
| with open( | |
| config.preprocess_text_config.train_path, mode="r", encoding="utf-8" | |
| ) as f: | |
| for line in f: | |
| speaker = line.split("|")[1] | |
| if speaker not in filelist_dict: | |
| filelist_dict[speaker] = [] | |
| filelist_dict[speaker].append(line.split("|")[0]) | |
| with Pool() as p: | |
| results = list( | |
| tqdm( | |
| p.imap(process_speaker, list(filelist_dict.keys())), | |
| total=len(filelist_dict), | |
| ) | |
| ) | |
| yml_result = {} | |
| for result in results: | |
| yml_result.update(result) | |
| with open( | |
| os.path.join(config.dataset_path, "emo_clustering/emo_clustering.yml"), | |
| "w", | |
| encoding="utf-8", | |
| ) as f: | |
| yaml.dump(yml_result, f) | |