Audio Enhancement with PyTorch: A Step-by-Step Guide

Audio Enhancement with PyTorch: A Step-by-Step Guide

This code demonstrates a practical application of PyTorch for audio enhancement. We'll walk through the process of training a generative model to reduce noise from audio signals.

Key Concepts:

• Audio Preprocessing: Slicing audio into manageable chunks and applying pre-emphasis for better noise reduction.
• Generator Network: A PyTorch model designed to generate clean audio from noisy input.
• Epoch Loading: Loading a pre-trained generator model for use in the enhancement process.
• Waveform Visualization: Comparing the original noisy audio and the enhanced version through waveform plots.

Code:

import argparse
import os
import wave

import numpy as np
import torch
import torch.nn as nn
from scipy.io import wavfile
from torch.autograd import Variable
from tqdm import tqdm
import matplotlib.pyplot as plt

from data_preprocess import slice_signal, window_size, sample_rate
from model import Generator
from utils import emphasis

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Test Single Audio Enhancement')
    parser.add_argument('--file_name', type=str, required=True, help='audio file name')
    parser.add_argument('--epoch_name', type=str, required=True, help='generator epoch name')

    opt = parser.parse_args()
    FILE_NAME = opt.file_name
    EPOCH_NAME = opt.epoch_name

    generator = Generator()
    generator.load_state_dict(torch.load('epochs/' + EPOCH_NAME, map_location='cpu'))
    if torch.cuda.is_available():
        generator.cuda()

    noisy_slices = slice_signal(FILE_NAME, window_size, 1, sample_rate)
    enhanced_speech = []
    for noisy_slice in tqdm(noisy_slices, desc='Generate enhanced audio'):
        z = nn.init.normal(torch.Tensor(1, 1024, 8))
        noisy_slice = torch.from_numpy(emphasis(noisy_slice[np.newaxis, np.newaxis, :])).type(torch.FloatTensor)
        if torch.cuda.is_available():
            noisy_slice, z = noisy_slice.cuda(), z.cuda()
        noisy_slice, z = Variable(noisy_slice), Variable(z)
        generated_speech = generator(noisy_slice, z).data.cpu().numpy()
        generated_speech = emphasis(generated_speech, emph_coeff=0.95, pre=False)
        generated_speech = generated_speech.reshape(-1)
        enhanced_speech.append(generated_speech)

    enhanced_speech = np.array(enhanced_speech).reshape(1, -1)
    file_name = os.path.join(os.path.dirname(FILE_NAME),
                             'enhanced_{}.wav'.format(os.path.basename(FILE_NAME).split('.')[0]))
    wavfile.write(file_name, sample_rate, enhanced_speech.T)


    # 读取音频文件
    w1 = wave.open('p232_036.wav', 'rb')
    w2 = wave.open('enhanced_p232_036.wav', 'rb')

    # 获取声道数、采样宽度和帧率
    channels1 = w1.getnchannels()
    channels2 = w2.getnchannels()
    sample_width1 = w1.getsampwidth()
    sample_width2 = w2.getsampwidth()
    framerate1 = w1.getframerate()
    framerate2 = w2.getframerate()

    # 读取音频数据
    frames1 = w1.readframes(-1)
    frames2 = w2.readframes(-1)

    # 将音频数据转换为numpy数组
    signal1 = np.frombuffer(frames1, dtype='int16')
    signal2 = np.frombuffer(frames2, dtype='int16')

    # 计算音频时长
    duration1 = len(signal1) / (framerate1 * channels1)
    duration2 = len(signal2) / (framerate2 * channels2)

    # 生成时间轴
    time1 = np.linspace(0, duration1, len(signal1))
    time2 = np.linspace(0, duration2, len(signal2))

    # 绘制波形图
    plt.figure()
    plt.subplot(2, 1, 1)
    plt.plot(time1, signal1)
    plt.title('Waveform of Audio 1')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.subplot(2, 1, 2)
    plt.plot(time2, signal2)
    plt.title('Waveform of Audio 2')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.show()

    # 关闭音频文件
    w1.close()
    w2.close()

Error Resolution:

The error wave.Error: unknown format: 3 usually arises when the audio file format is not recognized by the wave library. This can happen due to:

1. Incorrect File Format: Make sure the audio file you are loading is in WAV format (.wav).
2. Corrupted File: Verify that the file is not corrupted or damaged. Try downloading or generating the file again.
3. Encryption: Check if the file is encrypted. Decrypt the file before loading it.

If the issue persists despite these checks, consider using other libraries or tools specialized in audio handling, like librosa or pydub.