5.wer.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 随机抽取测试集计算WER与RTF (无GPU加速)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "import pathlib\n",
    "import pickle\n",
    "import random\n",
    "import json\n",
    "from functools import reduce\n",
    "\n",
    "# 禁用GPU\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "\n",
    "import librosa\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from scipy.fftpack import dct\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "keras = tf.keras\n",
    "ctc_decode, get_value, load_model = (\n",
    "    keras.backend.ctc_decode,\n",
    "    keras.backend.get_value,\n",
    "    keras.models.load_model,\n",
    ")\n",
    "\n",
    "# 音频/语音标注文件路径\n",
    "DS_PATH = \"data/\"\n",
    "# 模型文件路径\n",
    "FILES_PATH = \"output/\"\n",
    "\n",
    "# mfcc特征维数\n",
    "MFCC_VALUE = 32\n",
    "# FFT计算点数\n",
    "FFT = 512\n",
    "# 窗长\n",
    "WINLEN = 0.032\n",
    "\n",
    "# 音频文件路径\n",
    "data_path = sorted([str(p) for p in pathlib.Path(DS_PATH).glob(\"*.wav\")])\n",
    "# 语音标注文件路径\n",
    "label_path = sorted([str(p) for p in pathlib.Path(DS_PATH).glob(\"*.trn\")])\n",
    "\n",
    "# 音频时长\n",
    "SOUND_TIMES = []\n",
    "\n",
    "# 读取音频模型均值, 标准差\n",
    "with open(FILES_PATH + \"dataset/data_mfcc.pkl\", \"rb\") as file:\n",
    "    _, mfcc_mean, mfcc_std = pickle.load(file)\n",
    "    del _\n",
    "\n",
    "# 读取词库\n",
    "with open(FILES_PATH + \"dataset/words_vec.pkl\", \"rb\") as file:\n",
    "    char2id, id2char = pickle.load(file)\n",
    "\n",
    "\n",
    "def loadMfcc(filePath, mfccValue):\n",
    "    # 读取文件\n",
    "    data, sr = librosa.load(path=filePath, sr=None)\n",
    "    SOUND_TIMES.append(librosa.get_duration(y=data, sr=sr))\n",
    "\n",
    "    # 去除音频中所有的空白静默部分\n",
    "    y_split = librosa.effects.split(data, top_db=23)\n",
    "    y_split = np.array(\n",
    "        list(\n",
    "            reduce(\n",
    "                lambda x, y: np.concatenate((x, y)),\n",
    "                [data[x[0] : x[1]] for x in y_split],\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # 预加重\n",
    "    y_split = librosa.effects.preemphasis(y_split, coef=0.97)\n",
    "\n",
    "    # 提取MFCC特征\n",
    "    y_mfcc = librosa.feature.mfcc(\n",
    "        y=y_split,\n",
    "        sr=sr,\n",
    "        n_mfcc=mfccValue,\n",
    "        n_fft=512,\n",
    "        lifter=22,\n",
    "        hop_length=int(sr * 0.01),\n",
    "        win_length=int(sr * 0.025),\n",
    "    )\n",
    "    # 特征矩阵转置\n",
    "    y_mfcc = y_mfcc.transpose()\n",
    "    # 标准化\n",
    "    y_mfcc = (y_mfcc - mfcc_mean) / (mfcc_std + 1e-14)\n",
    "\n",
    "    return y_mfcc\n",
    "\n",
    "\n",
    "def ctc_decode_model(pred_audio, dict_list, model):\n",
    "    # 使用模型预测结果 (expand_dims用于增加维度匹配模型输入)\n",
    "    pred = model.predict(np.expand_dims(pred_audio, axis=0), verbose=0)\n",
    "    # 音频帧数\n",
    "    input_length = np.array((pred_audio.shape[1],))\n",
    "\n",
    "    # 解码\n",
    "    decode_res = ctc_decode(pred, input_length, greedy=True)\n",
    "    # 获取预测的序列数组, 筛选有效值 (> -1)\n",
    "    pred_index = get_value(decode_res[0][0])\n",
    "    pred_index = [item for item in pred_index[0] if item > -1]\n",
    "\n",
    "    # 使用词库将预测标号转换为文本\n",
    "    pred_text = \"\"\n",
    "    for index in pred_index:\n",
    "        pred_text += dict_list[index]\n",
    "\n",
    "    # 返回预测的文本及其标号\n",
    "    return pred_text, pred_index\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def levenshtein_distance(hypothesis: list, reference: list):\n",
    "    \"\"\"编辑距离\n",
    "    计算两个序列的levenshtein distance，可用于计算 WER/CER\n",
    "    参考资料：\n",
    "        https://www.cuelogic.com/blog/the-levenshtein-algorithm\n",
    "        https://martin-thoma.com/word-error-rate-calculation/\n",
    "\n",
    "    C: correct\n",
    "    W: wrong\n",
    "    I: insert\n",
    "    D: delete\n",
    "    S: substitution\n",
    "\n",
    "    :param hypothesis: 预测序列\n",
    "    :param reference: 真实序列\n",
    "    :return: 1: 错误操作，所需要的 S，D，I 操作的次数;\n",
    "             2: ref 与 hyp 的所有对齐下标\n",
    "             3: 返回 C、W、S、D、I 各自的数量\n",
    "    \"\"\"\n",
    "    len_hyp = len(hypothesis)\n",
    "    len_ref = len(reference)\n",
    "    cost_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int16)\n",
    "\n",
    "    # 记录所有的操作，0-equal；1-insertion；2-deletion；3-substitution\n",
    "    ops_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int8)\n",
    "\n",
    "    for i in range(len_hyp + 1):\n",
    "        cost_matrix[i][0] = i\n",
    "    for j in range(len_ref + 1):\n",
    "        cost_matrix[0][j] = j\n",
    "\n",
    "    # 生成 cost 矩阵和 operation矩阵，i:外层hyp，j:内层ref\n",
    "    for i in range(1, len_hyp + 1):\n",
    "        for j in range(1, len_ref + 1):\n",
    "            if hypothesis[i - 1] == reference[j - 1]:\n",
    "                cost_matrix[i][j] = cost_matrix[i - 1][j - 1]\n",
    "            else:\n",
    "                substitution = cost_matrix[i - 1][j - 1] + 1\n",
    "                insertion = cost_matrix[i - 1][j] + 1\n",
    "                deletion = cost_matrix[i][j - 1] + 1\n",
    "\n",
    "                # compare_val = [insertion, deletion, substitution]   # 优先级\n",
    "                compare_val = [substitution, insertion, deletion]  # 优先级\n",
    "\n",
    "                min_val = min(compare_val)\n",
    "                operation_idx = compare_val.index(min_val) + 1\n",
    "                cost_matrix[i][j] = min_val\n",
    "                ops_matrix[i][j] = operation_idx\n",
    "\n",
    "    match_idx = []  # 保存 hyp与ref 中所有对齐的元素下标\n",
    "    i = len_hyp\n",
    "    j = len_ref\n",
    "    nb_map = {\"N\": len_ref, \"C\": 0, \"W\": 0, \"I\": 0, \"D\": 0, \"S\": 0}\n",
    "    while i >= 0 or j >= 0:\n",
    "        i_idx = max(0, i)\n",
    "        j_idx = max(0, j)\n",
    "\n",
    "        if ops_matrix[i_idx][j_idx] == 0:  # correct\n",
    "            if i - 1 >= 0 and j - 1 >= 0:\n",
    "                match_idx.append((j - 1, i - 1))\n",
    "                nb_map[\"C\"] += 1\n",
    "\n",
    "            # 出边界后，这里仍然使用，应为第一行与第一列必然是全零的\n",
    "            i -= 1\n",
    "            j -= 1\n",
    "        # elif ops_matrix[i_idx][j_idx] == 1:   # insert\n",
    "        elif ops_matrix[i_idx][j_idx] == 2:  # insert\n",
    "            i -= 1\n",
    "            nb_map[\"I\"] += 1\n",
    "        # elif ops_matrix[i_idx][j_idx] == 2:   # delete\n",
    "        elif ops_matrix[i_idx][j_idx] == 3:  # delete\n",
    "            j -= 1\n",
    "            nb_map[\"D\"] += 1\n",
    "        # elif ops_matrix[i_idx][j_idx] == 3:   # substitute\n",
    "        elif ops_matrix[i_idx][j_idx] == 1:  # substitute\n",
    "            i -= 1\n",
    "            j -= 1\n",
    "            nb_map[\"S\"] += 1\n",
    "\n",
    "        # 出边界处理\n",
    "        if i < 0 and j >= 0:\n",
    "            nb_map[\"D\"] += 1\n",
    "        elif j < 0 and i >= 0:\n",
    "            nb_map[\"I\"] += 1\n",
    "\n",
    "    match_idx.reverse()\n",
    "    wrong_cnt = cost_matrix[len_hyp][len_ref]\n",
    "    nb_map[\"W\"] = wrong_cnt\n",
    "\n",
    "    return nb_map\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 随机读取测试集语音\n",
    "TEST_SIZE = 2496\n",
    "\n",
    "index_list = random.sample(range(0, len(data_path)), TEST_SIZE)\n",
    "sound_mfcc = []\n",
    "text_list = []\n",
    "\n",
    "\n",
    "for i in tqdm(index_list):\n",
    "    mfcc = loadMfcc(data_path[i], MFCC_VALUE)\n",
    "    sound_mfcc.append(mfcc)\n",
    "\n",
    "    # 读取语音对应文本\n",
    "    with open(label_path[i], \"r\", encoding=\"utf-8\") as file:\n",
    "        sound_text = file.readline().strip().replace(\" \", \"\")\n",
    "        text_list.append(sound_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载模型\n",
    "models_path = [str(x) for x in pathlib.Path(FILES_PATH + \"models/\").glob(\"**/*.h5\")]\n",
    "models = []\n",
    "\n",
    "for model in models_path:\n",
    "    models.append(load_model(model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wer_per_model = []\n",
    "rtf_per_model = []\n",
    "histories = {}\n",
    "\n",
    "for pre_model, path in tqdm(zip(models, models_path)):\n",
    "    wer_list = []\n",
    "    rtf_list = []\n",
    "    histories[path] = {\"data\": [], \"wer\": -1, \"rtf\": -1}\n",
    "\n",
    "    for sound, times, text in zip(sound_mfcc, SOUND_TIMES, text_list):\n",
    "        # 预测序列\n",
    "        start = time.time()\n",
    "        pred_text, _ = ctc_decode_model(sound, id2char, pre_model)\n",
    "        end = time.time() - start\n",
    "\n",
    "        # WER\n",
    "        nb_map = levenshtein_distance(hypothesis=list(pred_text), reference=list(text))\n",
    "        wer = ((nb_map[\"S\"] + nb_map[\"D\"] + nb_map[\"I\"]) / nb_map[\"N\"]) * 100\n",
    "        wer_list.append(wer)\n",
    "\n",
    "        # RTF\n",
    "        rtf = end / times\n",
    "        rtf_list.append(rtf)\n",
    "\n",
    "        # 保存记录\n",
    "        histories[path][\"data\"].append(\n",
    "            {\"reference\": text, \"hypothesis\": pred_text, \"wer\": wer, \"rtf\": rtf}\n",
    "        )\n",
    "\n",
    "    # 所有wer平均值\n",
    "    wer_per_model.append(np.mean(wer_list))\n",
    "    histories[path][\"wer\"] = np.mean(wer_list)\n",
    "    # 所有rtf平均值\n",
    "    rtf_per_model.append(np.mean(rtf_list))\n",
    "    histories[path][\"rtf\"] = np.mean(rtf_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, avg_wer, avg_rtf in zip(models_path, wer_per_model, rtf_per_model):\n",
    "    print(name, \"\\t\\t%.1f%%\" % avg_wer, \"\\t\\t%.4f\" % avg_rtf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存计算记录\n",
    "with open(\"output/wer_rtf.json\", \"w\", encoding=\"utf-8\") as file:\n",
    "    json.dump(histories, file, ensure_ascii=False)"
   ]
  }
 ],
 "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}