自动语音识别使用CTC

pip install jiwer

import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
from IPython import display
from jiwer import wer

data_url = "https://data.keithito.com/data/speech/LJSpeech-1.1.tar.bz2"
data_path = keras.utils.get_file("LJSpeech-1.1", data_url, untar=True)
wavs_path = data_path + "/wavs/"
metadata_path = data_path + "/metadata.csv"

# 读取元数据文件并解析
metadata_df = pd.read_csv(metadata_path, sep="|", header=None, quoting=3)
metadata_df.columns = ["file_name", "transcription", "normalized_transcription"]
metadata_df = metadata_df[["file_name", "normalized_transcription"]]
metadata_df = metadata_df.sample(frac=1).reset_index(drop=True)
metadata_df.head(3)

split = int(len(metadata_df) * 0.90)
df_train = metadata_df[:split]
df_val = metadata_df[split:]

print(f"训练集的大小: {len(df_train)}")
print(f"验证集的大小: {len(df_val)}")

# 转录中接受的字符集。
characters = [x for x in "abcdefghijklmnopqrstuvwxyz'?! "]
# 字符映射到整数
char_to_num = keras.layers.StringLookup(vocabulary=characters, oov_token="")
# 将整数映射回原始字符
num_to_char = keras.layers.StringLookup(
    vocabulary=char_to_num.get_vocabulary(), oov_token="", invert=True
)

print(
    f"词汇表为: {char_to_num.get_vocabulary()} "
    f"(大小 ={char_to_num.vocabulary_size()})"
)

# 一个整数标量张量。窗口长度（以样本为单位）。
frame_length = 256
# 一个整数标量张量。步进样本的数量。
frame_step = 160
# 一个整数标量张量。要应用的FFT大小。
# 如果未提供，则使用包含frame_length的最小2的幂。
fft_length = 384


def encode_single_sample(wav_file, label):
    ###########################################
    ##  处理音频
    ##########################################
    # 1. 读取wav文件
    file = tf.io.read_file(wavs_path + wav_file + ".wav")
    # 2. 解码wav文件
    audio, _ = tf.audio.decode_wav(file)
    audio = tf.squeeze(audio, axis=-1)
    # 3. 更改类型为浮点数
    audio = tf.cast(audio, tf.float32)
    # 4. 获取声谱图
    spectrogram = tf.signal.stft(
        audio, frame_length=frame_length, frame_step=frame_step, fft_length=fft_length
    )
    # 5. 我们只需要幅度，可以通过应用tf.abs获得
    spectrogram = tf.abs(spectrogram)
    spectrogram = tf.math.pow(spectrogram, 0.5)
    # 6. 归一化
    means = tf.math.reduce_mean(spectrogram, 1, keepdims=True)
    stddevs = tf.math.reduce_std(spectrogram, 1, keepdims=True)
    spectrogram = (spectrogram - means) / (stddevs + 1e-10)
    ###########################################
    ##  处理标签
    ##########################################
    # 7. 将标签转换为小写
    label = tf.strings.lower(label)
    # 8. 拆分标签
    label = tf.strings.unicode_split(label, input_encoding="UTF-8")
    # 9. 将标签中的字符映射到数字
    label = char_to_num(label)
    # 10. 返回字典，因为我们的模型期望两个输入
    return spectrogram, label

batch_size = 32
# 定义训练数据集
train_dataset = tf.data.Dataset.from_tensor_slices(
    (list(df_train["file_name"]), list(df_train["normalized_transcription"]))
)
train_dataset = (
    train_dataset.map(encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE)
    .padded_batch(batch_size)
    .prefetch(buffer_size=tf.data.AUTOTUNE)
)

# 定义验证数据集
validation_dataset = tf.data.Dataset.from_tensor_slices(
    (list(df_val["file_name"]), list(df_val["normalized_transcription"]))
)
validation_dataset = (
    validation_dataset.map(encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE)
    .padded_batch(batch_size)
    .prefetch(buffer_size=tf.data.AUTOTUNE)
)

fig = plt.figure(figsize=(8, 5))
for batch in train_dataset.take(1):
    spectrogram = batch[0][0].numpy()
    spectrogram = np.array([np.trim_zeros(x) for x in np.transpose(spectrogram)])
    label = batch[1][0]
    # 声谱图
    label = tf.strings.reduce_join(num_to_char(label)).numpy().decode("utf-8")
    ax = plt.subplot(2, 1, 1)
    ax.imshow(spectrogram, vmax=1)
    ax.set_title(label)
    ax.axis("off")
    # wav文件
    file = tf.io.read_file(wavs_path + list(df_train["file_name"])[0] + ".wav")
    audio, _ = tf.audio.decode_wav(file)
    audio = audio.numpy()
    ax = plt.subplot(2, 1, 2)
    plt.plot(audio)
    ax.set_title("信号波形")
    ax.set_xlim(0, len(audio))
    display.display(display.Audio(np.transpose(audio), rate=16000))
plt.show()

def CTCLoss(y_true, y_pred):
    # 计算训练时的损失值
    batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
    input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
    label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")

    input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")
    label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")

    loss = keras.backend.ctc_batch_cost(y_true, y_pred, input_length, label_length)
    return loss

def build_model(input_dim, output_dim, rnn_layers=5, rnn_units=128):
    """类似于DeepSpeech2的模型。"""
    # 模型输入
    input_spectrogram = layers.Input((None, input_dim), name="input")
    # 扩展维度以使用2D CNN。
    x = layers.Reshape((-1, input_dim, 1), name="expand_dim")(input_spectrogram)
    # 卷积层 1
    x = layers.Conv2D(
        filters=32,
        kernel_size=[11, 41],
        strides=[2, 2],
        padding="same",
        use_bias=False,
        name="conv_1",
    )(x)
    x = layers.BatchNormalization(name="conv_1_bn")(x)
    x = layers.ReLU(name="conv_1_relu")(x)
    # 卷积层 2
    x = layers.Conv2D(
        filters=32,
        kernel_size=[11, 21],
        strides=[1, 2],
        padding="same",
        use_bias=False,
        name="conv_2",
    )(x)
    x = layers.BatchNormalization(name="conv_2_bn")(x)
    x = layers.ReLU(name="conv_2_relu")(x)
    # 重新调整结果体积以输入RNN层
    x = layers.Reshape((-1, x.shape[-2] * x.shape[-1]))(x)
    # RNN层
    for i in range(1, rnn_layers + 1):
        recurrent = layers.GRU(
            units=rnn_units,
            activation="tanh",
            recurrent_activation="sigmoid",
            use_bias=True,
            return_sequences=True,
            reset_after=True,
            name=f"gru_{i}",
        )
        x = layers.Bidirectional(
            recurrent, name=f"bidirectional_{i}", merge_mode="concat"
        )(x)
        if i < rnn_layers:
            x = layers.Dropout(rate=0.5)(x)
    # 全连接层
    x = layers.Dense(units=rnn_units * 2, name="dense_1")(x)
    x = layers.ReLU(name="dense_1_relu")(x)
    x = layers.Dropout(rate=0.5)(x)
    # 分类层
    output = layers.Dense(units=output_dim + 1, activation="softmax")(x)
    # 模型
    model = keras.Model(input_spectrogram, output, name="DeepSpeech_2")
    # 优化器
    opt = keras.optimizers.Adam(learning_rate=1e-4)
    # 编译模型并返回
    model.compile(optimizer=opt, loss=CTCLoss)
    return model


# 获取模型
model = build_model(
    input_dim=fft_length // 2 + 1,
    output_dim=char_to_num.vocabulary_size(),
    rnn_units=512,
)
model.summary(line_length=110)