【Whisper】SpecAugmentationを実装してみた【ファインチューニング】

pip install numpy scipy

import numpy as np
import scipy.ndimage

def time_warping(spec, param_w=40):
    """時間方向のワーピング"""
    num_mel_channels, num_frames = spec.shape

    # num_frames が小さすぎる場合はワーピングせずにそのまま返す
    if num_frames <= 2 * param_w:
        return spec.copy()

    # 中心となる時間フレームを選択
    center = np.random.randint(param_w, num_frames - param_w)

    # 対象ピクセルの元の座標グリッド
    x, y = np.meshgrid(
        np.arange(num_frames), np.arange(num_mel_channels)
    )  # x: time, y: mel

    # ワープ量をガウス的に適用（時間方向にだけ変形）
    displacement = param_w * np.exp(-((x - center) ** 2) / (2 * (param_w**2)))

    # x座標だけずらす（縦方向（周波数）は固定）
    x_new = x + displacement
    y_new = y

    # 座標の形を合わせて補間
    coords = np.vstack([y_new.ravel(), x_new.ravel()])
    warped = scipy.ndimage.map_coordinates(spec, coords, order=3, mode="reflect")
    return warped.reshape(spec.shape)

 # num_frames が小さすぎる場合はワーピングせずにそのまま返す
    if num_frames <= 2 * param_w:
        return spec.copy()

def freq_masking(spec, param_f=27, num_masks=2, mask_value=0):
    """周波数方向のマスキング（mask_value を追加）"""
    spec = spec.copy()
    num_mel_channels = spec.shape[0]

    for _ in range(num_masks):
        if param_f < 1:
            continue

        f = np.random.randint(1, param_f + 1)  # f が 0 にならないように調整
        f0 = np.random.randint(0, num_mel_channels - f + 1)

        # マスキング適用（ゼロ埋め or ノイズ）
        if mask_value == 0:
            spec[f0 : f0 + f, :] = 0
        else:
            noise = np.random.normal(0, 1, (f, spec.shape[1])) * np.mean(spec)
            spec[f0 : f0 + f, :] = noise

    return spec

def time_masking(spec, param_p=0.2, param_T=40, num_masks=2, mask_value=0):
    spec = spec.copy()
    num_frames = spec.shape[1]

    # 全体でマスクしてよい最大量
    max_total_mask = int(param_p * num_frames)

    # 1回のマスク最大長さTをpに応じて制限
    adjusted_T = min(param_T, max_total_mask // num_masks) if num_masks > 0 else 0

    for _ in range(num_masks):
        if adjusted_T < 1:
            continue

        t = np.random.randint(1, adjusted_T + 1)
        t0 = np.random.randint(0, num_frames - t + 1)

        if mask_value == 0:
            spec[:, t0 : t0 + t] = 0
        else:
            noise = np.random.normal(0, 1, (spec.shape[0], t)) * np.mean(spec)
            spec[:, t0 : t0 + t] = noise

    return spec

def spec_augment(
    spec,
    param_w=40,
    param_f=15,
    param_T=40,
    param_p=0.2,
    num_freq_masks=2,
    num_time_masks=2,
    prob=0.5,
):
    if np.random.rand() < prob:
        spec = time_warping(spec, param_w=param_w)
        spec = freq_masking(spec, param_f=param_f, num_masks=num_freq_masks)
        spec = time_masking(
            spec, param_p=param_p, param_T=param_T, num_masks=num_time_masks
        )
    return spec

pip install librosa matplotlib

if __name__ == "__main__":
    import librosa
    import librosa.display
    import soundfile as sf
    import matplotlib.pyplot as plt


    # 音声ファイルの読み込み
    input_file = "test.mp3"
    y, sr = librosa.load(input_file, sr=None)

    # 短時間フーリエ変換（STFT）を取得
    stft_spec = librosa.stft(y, n_fft=1024, hop_length=512)
    stft_magnitude, stft_phase = np.abs(stft_spec), np.angle(stft_spec)
    stft_db = librosa.amplitude_to_db(stft_magnitude, ref=np.max)

    # Augmentationを適用
    augmented_stft_db = spec_augment(stft_db)

    # 変換後のスペクトログラムから音声を再構成
    augmented_stft_magnitude = librosa.db_to_amplitude(augmented_stft_db)
    augmented_stft = augmented_stft_magnitude * np.exp(
        1j * stft_phase
    )  # 位相情報を復元

    y_aug = librosa.istft(augmented_stft, hop_length=512)  # Griffin-Lim不要

    # 音量を元に戻す（正規化）
    y_aug = y_aug * (np.max(np.abs(y)) / np.max(np.abs(y_aug)))

    # 結果を保存（音声）
    sf.write("original.mp3", y, sr)
    sf.write("augmented.mp3", y_aug, sr)

    # スペクトログラムの可視化と保存
    fig, ax = plt.subplots(2, 1, figsize=(10, 8))
    librosa.display.specshow(stft_db, sr=sr, x_axis="time", y_axis="log", ax=ax[0])
    ax[0].set_title("Original STFT Spectrogram")
    librosa.display.specshow(
        augmented_stft_db, sr=sr, x_axis="time", y_axis="log", ax=ax[1]
    )
    ax[1].set_title("Augmented STFT Spectrogram")
    plt.tight_layout()
    plt.savefig("spectrogram_comparison.png")  # 画像として保存