Upload classify.py

classify.py

@@ -0,0 +1,180 @@
+from sklearn.cluster import *
+import os
+import numpy as np
+from config import config
+import yaml
+import argparse
+import shutil
+
+def ensure_dir(directory):
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-a", "--algorithm", default="k", help="choose algorithm", type=str)
+    parser.add_argument("-n", "--num_clusters", default=4, help="number of clusters", type=int)
+    parser.add_argument("-r", "--range", default=4, help="number of files in a class", type=int)
+    args = parser.parse_args()
+
+    filelist_dict = {}
+    yml_result = {}
+    base_dir = "D:/Vits2/Bert-VITS2/Data/BanGDream/filelists"
+    output_dir = "D:/Vits2/classifedSample"
+
+    with open(os.path.join(base_dir, "Mygo.list"), mode="r", encoding="utf-8") as f:
+        embs = []
+        wavnames = []
+        for line in f:
+            parts = line.strip().split("|")
+            speaker = parts[1]  # 假设 speaker 是第二个部分
+            filepath = parts[0]  # 假设 filepath 是第一个部分
+            # ... 其余部分可以根据需要使用
+
+            if speaker not in filelist_dict:
+                filelist_dict[speaker] = []
+                yml_result[speaker] = {}
+            filelist_dict[speaker].append(filepath)
+
+    for speaker in filelist_dict:
+        print("\nspeaker: " + speaker)
+
+        embs = []
+        wavnames = []
+
+        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(file)
+            except Exception as e:
+                print(e)
+
+        if embs:
+            n_clusters = args.num_clusters
+            x = np.concatenate(embs, axis=0)
+            x = np.squeeze(x)
+
+            if args.algorithm == "b":
+                model = Birch(n_clusters=n_clusters, threshold=0.2)
+            elif args.algorithm == "s":
+                model = SpectralClustering(n_clusters=n_clusters)
+            elif args.algorithm == "a":
+                model = AgglomerativeClustering(n_clusters=n_clusters)
+            else:
+                model = KMeans(n_clusters=n_clusters, random_state=10)
+
+            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)
+
+            for i in range(y_predict.max() + 1):
+                print("类别:", i, "本类中样本数量:", len(classes[i]))
+                yml_result[speaker][f"class{i}"] = []
+                class_dir = os.path.join(output_dir, speaker, f"class{i}")
+
+                num_samples_in_class = len(classes[i])
+                for j in range(min(args.range, num_samples_in_class)):
+                    wav_file = classes[i][j]
+                    print(wav_file)
+
+                    # 复制文件到新目录
+                    ensure_dir(class_dir)
+                    shutil.copy(os.path.join(base_dir, wav_file), os.path.join(class_dir, os.path.basename(wav_file)))
+
+                    yml_result[speaker][f"class{i}"].append(wav_file)
+
+    with open(os.path.join(base_dir, "emo_clustering.yml"), "w", encoding="utf-8") as f:
+        yaml.dump(yml_result, f)
+
+'''
+from sklearn.cluster import *
+import os
+import numpy as np
+from config import config
+import yaml
+import argparse
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-a", "--algorithm", default="s", help="choose algorithm", type=str
+    )
+    parser.add_argument(
+        "-n", "--num_clusters", default=3, help="number of clusters", type=int
+    )
+    parser.add_argument(
+        "-r", "--range", default=4, help="number of files in a class", type=int
+    )
+    args = parser.parse_args()
+    filelist_dict = {}
+    yml_result = {}
+    with open(
+        "D:/Vits2/Bert-VITS2/Data/BanGDream/filelists/Mygo.list", mode="r", encoding="utf-8"
+    ) as f:
+        embs = []
+        wavnames = []
+        for line in f:
+            speaker = line.split("|")[1]
+            if speaker not in filelist_dict:
+                filelist_dict[speaker] = []
+                yml_result[speaker] = {}
+            filelist_dict[speaker].append(line.split("|")[0])
+    #print(filelist_dict)
+
+    for speaker in filelist_dict:
+        print("\nspeaker: " + speaker)
+
+        # 清空 embs 和 wavnames 列表
+        embs = []
+        wavnames = []
+
+        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)
+
+        if embs:
+        # 聚类算法类的数量
+            n_clusters = args.num_clusters
+            x = np.concatenate(embs, axis=0)
+            x = np.squeeze(x)
+            # 聚类算法类的数量
+            n_clusters = args.num_clusters
+            if args.algorithm == "b":
+                model = Birch(n_clusters=n_clusters, threshold=0.2)
+            elif args.algorithm == "s":
+                model = SpectralClustering(n_clusters=n_clusters)
+            elif args.algorithm == "a":
+                model = AgglomerativeClustering(n_clusters=n_clusters)
+            else:
+                model = KMeans(n_clusters=n_clusters, random_state=10)
+            # 可以自行尝试各种不同的聚类算法
+            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)
+
+            for i in range(y_predict.max() + 1):
+                print("类别:", i, "本类中样本数量:", len(classes[i]))
+                yml_result[speaker][f"class{i}"] = []
+
+                # 修正：确保不会尝试访问超出范围的元素
+                num_samples_in_class = len(classes[i])
+                for j in range(min(args.range, num_samples_in_class)):
+                    print(classes[i][j])
+                    yml_result[speaker][f"class{i}"].append(classes[i][j])
+    with open(
+        os.path.join('D:/Vits2/Bert-VITS2/Data/BanGDream', "emo_clustering.yml"), "w", encoding="utf-8"
+    ) as f:
+        yaml.dump(yml_result, f)
+        '''