{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# スタイルベクトルをもっと詳しく作りたい人向け\n",
        "\n",
        "JVNVコーパスの例を使いながらクラスタリングをいろいろいじったり、また正解ラベルからベクトルを作りたい人向けです。\n",
        "もっといろいろ足したりいろんなスタイルベクトルを作って遊べると思います。\n",
        "ある程度慣れている人向けです。\n",
        "\n",
        "例では[JVNVコーパス](https://sites.google.com/site/shinnosuketakamichi/research-topics/jvnv_corpus)のjvnv-M1を使います。\n",
        "\n",
        "## そもそもスタイルベクトルとは\n",
        "[この話者識別モデル](https://huggingface.co/pyannote/wespeaker-voxceleb-resnet34-LM)を使って生成された、1つの音声ファイルにつき256次元のベクトルです。話者識別用のものですが、感情や声音の特徴も含まれているので、スタイルベクトルとして使えます。\n",
        "\n",
        "このStyle-Bert-VITS2では、この256次元のベクトルをエンコーダに注入して学習しているので、推論時にそのベクトルを入れてあげる必要があります。ある感情を強く表していると思われるベクトルを入れると、その感情を強く表現した音声が生成される、という仕組みです。\n",
        "\n",
        "もともとが話者識別用なので、「この感情はこのベクトル」のような普遍的なスタイルベクトルは使えません。なのでこのように、いちいちデータセットごとにベクトルを作る必要があります。\n",
        "\n",
        "## モデルを使うために必要なもの\n",
        "- `model_assets/{model_name}/model_name.safetensors`: 学習の結果出力されるモデルファイル。これは自動的にこの場所に置かれ、スタイルベクトルとは全く独立。\n",
        "- `model_assets/{model_name}/style_vectors.npy`: スタイルベクトルのnumpyファイル。ベクトルをいくつかいれる。\n",
        "- `model_assets/{model_name}/config.json`: モデルの設定ファイル（学習前準備で自動的に生成されるはず）。これの以下の項目を設定。\n",
        "```json\n",
        "{\n",
        "    \"data\": {\n",
        "        \"num_styles\": 4, // スタイルベクトルの数\n",
        "        \"style2id\": { // スタイルベクトルの名前とidの対応\n",
        "            \"Neutral\": 0,\n",
        "            \"Angry\": 1,\n",
        "            \"Happy\": 2,\n",
        "            \"Sad\": 3\n",
        "        }\n",
        "    }\n",
        "}\n",
        "```\n",
        "ここでidは0から始まる整数で、スタイルベクトルのnumpyファイルの何番目のベクトルかを指定します。最初のNeutralは含めたほうがよさそうで、WebUIや下では全スタイルベクトルの平均を入れています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "\n",
        "wav_dir = \"Data/jvnv-M1/wavs\"\n",
        "\n",
        "embs = []\n",
        "names= []\n",
        "for file in os.listdir(wav_dir):\n",
        "    if file.endswith(\".npy\"):\n",
        "        xvec = np.load(os.path.join(wav_dir, file))\n",
        "        embs.append(np.expand_dims(xvec, axis=0))\n",
        "        names.append(file)\n",
        "\n",
        "x = np.concatenate(embs, axis=0)\n",
        "x = np.squeeze(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 36548 characters]",
            "text/plain": [
              "<Figure size 700x700 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# TSNEで可視化\n",
        "from sklearn.manifold import TSNE\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "# 特徴量の関係か、コサイン距離が良さげ\n",
        "tsne = TSNE(n_components=2, random_state=42, metric=\"cosine\")\n",
        "\n",
        "x_tsne = tsne.fit_transform(x)\n",
        "\n",
        "plt.figure(figsize=(7, 7))\n",
        "plt.scatter(x_tsne[:, 0], x_tsne[:, 1])\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "from sklearn.cluster import KMeans, AgglomerativeClustering\n",
        "\n",
        "method = \"k\"\n",
        "n_clusters = 5\n",
        "\n",
        "if method == \"k\":\n",
        "    model = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)\n",
        "elif method == \"a\":\n",
        "    model = AgglomerativeClustering(n_clusters=n_clusters)\n",
        "\n",
        "y_predict = model.fit_predict(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "上は直接生のベクトルでクラスタリングしてますが、次のようにt-SNEで図のように2次元に削減したものをクラスタリングしたほうが、識別がきれいに分かれる、ことが多いような気がします、が詳しくは分かりません:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "y_predict = model.fit_predict(x_tsne)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "centroids = []\n",
        "for i in range(n_clusters):\n",
        "    centroids.append(x[y_predict == i].mean(axis=0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 104320 characters]",
            "text/plain": [
              "<Figure size 700x700 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# TSNEで可視化\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.figure(figsize=(7, 7))\n",
        "plt.scatter(x_tsne[:, 0], x_tsne[:, 1], c=y_predict)\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from scipy.spatial.distance import cdist\n",
        "from IPython.display import Audio, display\n",
        "\n",
        "# 各クラスターのセントロイドに最も近い点に対応するファイル名を取得\n",
        "closest_files = []\n",
        "for center_idx in range(len(centroids)):\n",
        "    closest_idx = np.argmin(\n",
        "        cdist(centroids[center_idx : center_idx + 1], x, metric=\"cosine\")\n",
        "    )\n",
        "    closest_files.append(names[closest_idx])\n",
        "\n",
        "# 対応する音声ファイルをJupyterノートブック上で再生\n",
        "for file_name in closest_files:\n",
        "    wav_path = os.path.join(wav_dir, file_name.replace(\".npy\", \"\"))\n",
        "    if os.path.exists(wav_path):\n",
        "        print(wav_path)\n",
        "        display(Audio(wav_path))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {},
      "outputs": [],
      "source": [
        "# meanとcentroidを保存\n",
        "mean = x.mean(axis=0)\n",
        "save_vectors = np.vstack([mean, centroids])\n",
        "os.makedirs(\"model_assets/jvnv-M1\", exist_ok=True)\n",
        "np.save(\"model_assets/jvnv-M1/style_vectors.npy\", save_vectors)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "# JVNVコーパスの場合は正解はファイル名の最初の方から分かる。それを使って正解ラベルを作成\n",
        "y_true = [name[3:6] for name in names]\n",
        "\n",
        "for i, y in enumerate(y_true):\n",
        "    if y == \"ang\":\n",
        "        y_true[i] = 0\n",
        "    elif y == \"dis\":\n",
        "        y_true[i] = 1\n",
        "    elif y == \"fea\":\n",
        "        y_true[i] = 2\n",
        "    elif y == \"hap\":\n",
        "        y_true[i] = 3\n",
        "    elif y == \"sad\":\n",
        "        y_true[i] = 4\n",
        "    elif y == \"sur\":\n",
        "        y_true[i] = 5\n",
        "    else:\n",
        "        print(y)\n",
        "        raise ValueError"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 正解のセントロイドを計算\n",
        "y_true = np.array(y_true)\n",
        "true_centroids = []\n",
        "for i in range(6):\n",
        "    true_centroids.append(np.mean(x[y_true == i], axis=0))\n",
        "\n",
        "true_centroids = np.array(true_centroids)\n",
        "\n",
        "# すべてのベクトルの平均を計算\n",
        "mean = np.mean(x, axis=0)\n",
        "\n",
        "# 保存\n",
        "os.makedirs(\"model_assets/jvnv-M1\", exist_ok=True)\n",
        "np.save(\"model_assets/jvnv-M1/style_vectors.npy\", np.vstack([mean, true_centroids]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 47576 characters]",
            "text/plain": [
              "<Figure size 700x700 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# TSNEで正解ラベルを可視化\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "cmap = plt.get_cmap(\"tab10\")\n",
        "\n",
        "true_label_dict = {0: \"ang\", 1: \"dis\", 2: \"fea\", 3: \"hap\", 4: \"sad\", 5: \"sur\"}\n",
        "\n",
        "plt.figure(figsize=(7, 7))\n",
        "for i in true_label_dict:\n",
        "    plt.scatter(\n",
        "        x_tsne[y_true == i, 0],\n",
        "        x_tsne[y_true == i, 1],\n",
        "        color=cmap(i),\n",
        "        label=f\"{true_label_dict[i]}\",\n",
        "    )\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
