From 900ed8aad4baefefad588502c1f031043c8665cb Mon Sep 17 00:00:00 2001
From: KakaruHayate <kakaru.hk@hotmail.com>
Date: Wed, 24 Jan 2024 19:09:41 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8A=BD=E5=8D=A1=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

抽卡功能，支持使用colorsplitter生成的CSV分类文件

https://github.com/KakaruHayate/ColorSplitter
---
 GPT_SoVITS/inference_webui.py | 35 +++++++++++++++++++++++++++++++++++
 requirements.txt              |  1 +
 2 files changed, 36 insertions(+)

diff --git a/GPT_SoVITS/inference_webui.py b/GPT_SoVITS/inference_webui.py
index c66305b0..4035915b 100644
--- a/GPT_SoVITS/inference_webui.py
+++ b/GPT_SoVITS/inference_webui.py
@@ -33,6 +33,8 @@ from time import time as ttime
 from module.mel_processing import spectrogram_torch
 from my_utils import load_audio
 from tools.i18n.i18n import I18nAuto
+import random
+import pandas as pd
 i18n = I18nAuto()
 
 device = "cuda"
@@ -418,6 +420,7 @@ def change_choices():
     SoVITS_names, GPT_names = get_weights_names()
     return {"choices": sorted(SoVITS_names,key=custom_sort_key), "__type__": "update"}, {"choices": sorted(GPT_names,key=custom_sort_key), "__type__": "update"}
 
+
 pretrained_sovits_name="GPT_SoVITS/pretrained_models/s2G488k.pth"
 pretrained_gpt_name="GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt"
 SoVITS_weight_root="SoVITS_weights"
@@ -434,6 +437,29 @@ def get_weights_names():
     return SoVITS_names,GPT_names
 SoVITS_names,GPT_names = get_weights_names()
 
+
+def select_random_file(folder_path):
+    files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
+    return os.path.join(folder_path, random.choice(files))
+
+
+def get_labels(csv_path):
+    if csv_path:
+        df = pd.read_csv(csv_path)
+        return {"choices": df['clust'].unique().tolist(), "__type__": "update"}
+    else:
+        return {"choices": [], "__type__": "update"}
+
+
+def select_random_file_from_csv(csv_path, label, folder_path):
+    if csv_path:
+        df = pd.read_csv(csv_path)
+        files = df[df['clust'] == label]['filename'].apply(os.path.basename).tolist()
+        return os.path.join(folder_path, random.choice(files))
+    else:
+        return select_random_file(folder_path)
+
+
 with gr.Blocks(title="GPT-SoVITS WebUI") as app:
     gr.Markdown(
         value=i18n("本软件以MIT协议开源, 作者不对软件具备任何控制力, 使用软件者、传播软件导出的声音者自负全责. <br>如不认可该条款, 则不能使用或引用软件包内任何代码和文件. 详见根目录<b>LICENSE</b>.")
@@ -447,9 +473,18 @@ with gr.Blocks(title="GPT-SoVITS WebUI") as app:
             refresh_button.click(fn=change_choices, inputs=[], outputs=[SoVITS_dropdown, GPT_dropdown])
             SoVITS_dropdown.change(change_sovits_weights,[SoVITS_dropdown],[])
             GPT_dropdown.change(change_gpt_weights,[GPT_dropdown],[])
+        gr.Markdown(value="随机抽取参考音频，可以导入csv文件限定抽取范围，csv第一列为filename，第二列为clust（类别标签）")
+        with gr.Row():
+            folder_path = gr.Textbox(label="请输入音频文件夹路径", value="")
+            csv_path = gr.Textbox(label="请输入CSV文件路径（可选）", value="")
+            label = gr.Dropdown(label="请选择一个标签（可选）", choices=[], value="")
+            update_labels_button = gr.Button("更新标签", variant="primary")
+            update_labels_button.click(get_labels, [csv_path], [label])
+            random_file_button = gr.Button("随机选择音频文件", variant="primary")
         gr.Markdown(value=i18n("*请上传并填写参考信息"))
         with gr.Row():
             inp_ref = gr.Audio(label=i18n("请上传参考音频"), type="filepath")
+            random_file_button.click(select_random_file_from_csv, [csv_path, label, folder_path], [inp_ref])
             prompt_text = gr.Textbox(label=i18n("参考音频的文本"), value="")
             prompt_language = gr.Dropdown(
                 label=i18n("参考音频的语种"),choices=[i18n("中文"),i18n("英文"),i18n("日文")],value=i18n("中文")
diff --git a/requirements.txt b/requirements.txt
index a8e72ea9..006c1f6e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -21,3 +21,4 @@ chardet
 PyYAML
 psutil
 jieba_fast
+pandas
\ No newline at end of file