PaddleSpeech/third_party/nnAudio/tests/test_spectrogram.py

import pytest
import librosa
import torch
import matplotlib.pyplot as plt
from scipy.signal import chirp, sweep_poly
from nnAudio.Spectrogram import *
from parameters import *

gpu_idx=0

# librosa example audio for testing
example_y, example_sr = librosa.load(librosa.util.example_audio_file())


@pytest.mark.parametrize("n_fft, hop_length, window", stft_parameters)  
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_inverse2(n_fft, hop_length, window, device):
    x = torch.tensor(example_y,device=device)
    stft = STFT(n_fft=n_fft, hop_length=hop_length, window=window).to(device)
    istft = iSTFT(n_fft=n_fft, hop_length=hop_length, window=window).to(device)
    X = stft(x.unsqueeze(0), output_format="Complex")
    x_recon = istft(X, length=x.shape[0], onesided=True).squeeze()
    assert np.allclose(x.cpu(), x_recon.cpu(), rtol=1e-5, atol=1e-3)    

@pytest.mark.parametrize("n_fft, hop_length, window", stft_parameters)
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_inverse(n_fft, hop_length, window, device):
    x = torch.tensor(example_y, device=device)
    stft = STFT(n_fft=n_fft, hop_length=hop_length, window=window, iSTFT=True).to(device)
    X = stft(x.unsqueeze(0), output_format="Complex")
    x_recon = stft.inverse(X, length=x.shape[0]).squeeze()
    assert np.allclose(x.cpu(), x_recon.cpu(), rtol=1e-3, atol=1)
    

    
# @pytest.mark.parametrize("n_fft, hop_length, window", stft_parameters)

# def test_inverse_GPU(n_fft, hop_length, window):
#     x = torch.tensor(example_y,device=f'cuda:{gpu_idx}')
#     stft = STFT(n_fft=n_fft, hop_length=hop_length, window=window, device=f'cuda:{gpu_idx}')
#     X = stft(x.unsqueeze(0), output_format="Complex")
#     x_recon = stft.inverse(X, num_samples=x.shape[0]).squeeze()
#     assert np.allclose(x.cpu(), x_recon.cpu(), rtol=1e-3, atol=1)


@pytest.mark.parametrize("n_fft, hop_length, window", stft_parameters)
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_stft_complex(n_fft, hop_length, window, device):
    x = example_y
    stft = STFT(n_fft=n_fft, hop_length=hop_length, window=window).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0), output_format="Complex")
    X_real, X_imag = X[:, :, :, 0].squeeze(), X[:, :, :, 1].squeeze()
    X_librosa = librosa.stft(x, n_fft=n_fft, hop_length=hop_length, window=window)
    real_diff, imag_diff = np.allclose(X_real.cpu(), X_librosa.real, rtol=1e-3, atol=1e-3), \
                            np.allclose(X_imag.cpu(), X_librosa.imag, rtol=1e-3, atol=1e-3)
    
    assert real_diff and imag_diff 
    
# @pytest.mark.parametrize("n_fft, hop_length, window", stft_parameters)    
# def test_stft_complex_GPU(n_fft, hop_length, window):
#     x = example_y
#     stft = STFT(n_fft=n_fft, hop_length=hop_length, window=window, device=f'cuda:{gpu_idx}')
#     X = stft(torch.tensor(x,device=f'cuda:{gpu_idx}').unsqueeze(0), output_format="Complex")
#     X_real, X_imag = X[:, :, :, 0].squeeze().detach().cpu(), X[:, :, :, 1].squeeze().detach().cpu()
#     X_librosa = librosa.stft(x, n_fft=n_fft, hop_length=hop_length, window=window)
#     real_diff, imag_diff = np.allclose(X_real, X_librosa.real, rtol=1e-3, atol=1e-3), \
#                             np.allclose(X_imag, X_librosa.imag, rtol=1e-3, atol=1e-3)
    
#     assert real_diff and imag_diff        
    
@pytest.mark.parametrize("n_fft, win_length, hop_length", stft_with_win_parameters) 
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_stft_complex_winlength(n_fft, win_length, hop_length, device):
    x = example_y
    stft = STFT(n_fft=n_fft, win_length=win_length, hop_length=hop_length).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0), output_format="Complex")
    X_real, X_imag = X[:, :, :, 0].squeeze(), X[:, :, :, 1].squeeze()
    X_librosa = librosa.stft(x, n_fft=n_fft, win_length=win_length, hop_length=hop_length)
    real_diff, imag_diff = np.allclose(X_real.cpu(), X_librosa.real, rtol=1e-3, atol=1e-3), \
                            np.allclose(X_imag.cpu(), X_librosa.imag, rtol=1e-3, atol=1e-3)
    assert real_diff and imag_diff    
              
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_stft_magnitude(device):
    x = example_y
    stft = STFT(n_fft=2048, hop_length=512).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0), output_format="Magnitude").squeeze()
    X_librosa, _ = librosa.core.magphase(librosa.stft(x, n_fft=2048, hop_length=512))
    assert np.allclose(X.cpu(), X_librosa, rtol=1e-3, atol=1e-3)

@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_stft_phase(device):
    x = example_y
    stft = STFT(n_fft=2048, hop_length=512).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0), output_format="Phase")
    X_real, X_imag = torch.cos(X).squeeze(), torch.sin(X).squeeze()
    _, X_librosa = librosa.core.magphase(librosa.stft(x, n_fft=2048, hop_length=512))

    real_diff, imag_diff = np.mean(np.abs(X_real.cpu().numpy() - X_librosa.real)), \
                            np.mean(np.abs(X_imag.cpu().numpy() - X_librosa.imag))

    # I find that np.allclose is too strict for allowing phase to be similar to librosa.
    # Hence for phase we use average element-wise distance as the test metric.
    assert real_diff < 2e-4 and imag_diff < 2e-4

@pytest.mark.parametrize("n_fft, win_length", mel_win_parameters)  
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_mel_spectrogram(n_fft, win_length, device):
    x = example_y
    melspec = MelSpectrogram(n_fft=n_fft, win_length=win_length, hop_length=512).to(device)
    X = melspec(torch.tensor(x, device=device).unsqueeze(0)).squeeze()
    X_librosa = librosa.feature.melspectrogram(x, n_fft=n_fft, win_length=win_length, hop_length=512)
    assert np.allclose(X.cpu(), X_librosa, rtol=1e-3, atol=1e-3)
    
    
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_cqt_1992(device):
    # Log sweep case
    fs = 44100
    t = 1
    f0 = 55
    f1 = 22050
    s = np.linspace(0, t, fs*t)
    x = chirp(s, f0, 1, f1, method='logarithmic')
    x = x.astype(dtype=np.float32)

    # Magnitude
    stft = CQT1992(sr=fs, fmin=220, output_format="Magnitude",
                     n_bins=80, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))


    # Complex
    stft = CQT1992(sr=fs, fmin=220, output_format="Complex",
                     n_bins=80, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))

    # Phase
    stft = CQT1992(sr=fs, fmin=220, output_format="Phase",
                     n_bins=160, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    
    assert True
    
@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_cqt_2010(device):
    # Log sweep case
    fs = 44100
    t = 1
    f0 = 55
    f1 = 22050
    s = np.linspace(0, t, fs*t)
    x = chirp(s, f0, 1, f1, method='logarithmic')
    x = x.astype(dtype=np.float32)

    # Magnitude
    stft = CQT2010(sr=fs, fmin=110, output_format="Magnitude",
                     n_bins=160, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))     

    # Complex
    stft = CQT2010(sr=fs, fmin=110, output_format="Complex",
                     n_bins=160, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    
    # Phase
    stft = CQT2010(sr=fs, fmin=110, output_format="Phase",
                     n_bins=160, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))    
    assert True   

@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_cqt_1992_v2_log(device):
    # Log sweep case
    fs = 44100
    t = 1
    f0 = 55
    f1 = 22050
    s = np.linspace(0, t, fs*t)
    x = chirp(s, f0, 1, f1, method='logarithmic')
    x = x.astype(dtype=np.float32)

    # Magnitude
    stft = CQT1992v2(sr=fs, fmin=55, output_format="Magnitude",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    ground_truth = np.load("tests/ground-truths/log-sweep-cqt-1992-mag-ground-truth.npy")
    X = torch.log(X + 1e-5)
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Complex
    stft = CQT1992v2(sr=fs, fmin=55, output_format="Complex",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    ground_truth = np.load("tests/ground-truths/log-sweep-cqt-1992-complex-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Phase
    stft = CQT1992v2(sr=fs, fmin=55, output_format="Phase",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    ground_truth = np.load("tests/ground-truths/log-sweep-cqt-1992-phase-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_cqt_1992_v2_linear(device):
    # Linear sweep case
    fs = 44100
    t = 1
    f0 = 55
    f1 = 22050
    s = np.linspace(0, t, fs*t)
    x = chirp(s, f0, 1, f1, method='linear')
    x = x.astype(dtype=np.float32)

    # Magnitude
    stft = CQT1992v2(sr=fs, fmin=55, output_format="Magnitude",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    ground_truth = np.load("tests/ground-truths/linear-sweep-cqt-1992-mag-ground-truth.npy")
    X = torch.log(X + 1e-5)
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Complex
    stft = CQT1992v2(sr=fs, fmin=55, output_format="Complex",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    ground_truth = np.load("tests/ground-truths/linear-sweep-cqt-1992-complex-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Phase
    stft = CQT1992v2(sr=fs, fmin=55, output_format="Phase",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    ground_truth = np.load("tests/ground-truths/linear-sweep-cqt-1992-phase-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_cqt_2010_v2_log(device):
    # Log sweep case
    fs = 44100
    t = 1
    f0 = 55
    f1 = 22050
    s = np.linspace(0, t, fs*t)
    x = chirp(s, f0, 1, f1, method='logarithmic')
    x = x.astype(dtype=np.float32)

    # Magnitude
    stft = CQT2010v2(sr=fs, fmin=55, output_format="Magnitude",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))     
    X = torch.log(X + 1e-2)
#     np.save("tests/ground-truths/log-sweep-cqt-2010-mag-ground-truth", X.cpu())
    ground_truth = np.load("tests/ground-truths/log-sweep-cqt-2010-mag-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Complex
    stft = CQT2010v2(sr=fs, fmin=55, output_format="Complex",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
#     np.save("tests/ground-truths/log-sweep-cqt-2010-complex-ground-truth", X.cpu())         
    ground_truth = np.load("tests/ground-truths/log-sweep-cqt-2010-complex-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

#     # Phase
#     stft = CQT2010v2(sr=fs, fmin=55, device=device, output_format="Phase",
#                      n_bins=207, bins_per_octave=24)
#     X = stft(torch.tensor(x, device=device).unsqueeze(0))
# #     np.save("tests/ground-truths/log-sweep-cqt-2010-phase-ground-truth", X.cpu())      
#     ground_truth = np.load("tests/ground-truths/log-sweep-cqt-2010-phase-ground-truth.npy")
#     assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_cqt_2010_v2_linear(device):
    # Linear sweep case
    fs = 44100
    t = 1
    f0 = 55
    f1 = 22050
    s = np.linspace(0, t, fs*t)
    x = chirp(s, f0, 1, f1, method='linear')
    x = x.astype(dtype=np.float32)

    # Magnitude
    stft = CQT2010v2(sr=fs, fmin=55, output_format="Magnitude",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
    X = torch.log(X + 1e-2)
#     np.save("tests/ground-truths/linear-sweep-cqt-2010-mag-ground-truth", X.cpu()) 
    ground_truth = np.load("tests/ground-truths/linear-sweep-cqt-2010-mag-ground-truth.npy")    
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Complex
    stft = CQT2010v2(sr=fs, fmin=55, output_format="Complex",
                     n_bins=207, bins_per_octave=24).to(device)
    X = stft(torch.tensor(x, device=device).unsqueeze(0))
#     np.save("tests/ground-truths/linear-sweep-cqt-2010-complex-ground-truth", X.cpu())
    ground_truth = np.load("tests/ground-truths/linear-sweep-cqt-2010-complex-ground-truth.npy")
    assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

    # Phase
#     stft = CQT2010v2(sr=fs, fmin=55, device=device, output_format="Phase",
#                      n_bins=207, bins_per_octave=24)
#     X = stft(torch.tensor(x, device=device).unsqueeze(0))
# #     np.save("tests/ground-truths/linear-sweep-cqt-2010-phase-ground-truth", X.cpu())
#     ground_truth = np.load("tests/ground-truths/linear-sweep-cqt-2010-phase-ground-truth.npy")
#     assert np.allclose(X.cpu(), ground_truth, rtol=1e-3, atol=1e-3)

@pytest.mark.parametrize("device", ['cpu', f'cuda:{gpu_idx}'])
def test_mfcc(device):
    x = example_y
    mfcc = MFCC(sr=example_sr).to(device)
    X = mfcc(torch.tensor(x, device=device).unsqueeze(0)).squeeze()
    X_librosa = librosa.feature.mfcc(x, sr=example_sr)
    assert np.allclose(X.cpu(), X_librosa, rtol=1e-3, atol=1e-3)
    

x = torch.randn((4,44100)) # Create a batch of input for the following Data.Parallel test

@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])
def test_STFT_Parallel(device):
    spec_layer = STFT(hop_length=512, n_fft=2048, window='hann', 
                                  freq_scale='no',
                                  output_format='Complex').to(device)
    inverse_spec_layer = iSTFT(hop_length=512, n_fft=2048, window='hann', 
                                  freq_scale='no').to(device)    
    
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    inverse_spec_layer_parallel = torch.nn.DataParallel(inverse_spec_layer)
    spec = spec_layer_parallel(x)
    x_recon = inverse_spec_layer_parallel(spec, onesided=True, length=x.shape[-1])
    
    assert np.allclose(x_recon.detach().cpu(), x.detach().cpu(), rtol=1e-3, atol=1e-3)

@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])   
def test_MelSpectrogram_Parallel(device):
    spec_layer = MelSpectrogram(sr=22050, n_fft=2048, n_mels=128, hop_length=512,
                                            window='hann', center=True, pad_mode='reflect', 
                                            power=2.0, htk=False, fmin=0.0, fmax=None, norm=1, 
                                            verbose=True).to(device)
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    spec = spec_layer_parallel(x)

@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])    
def test_MFCC_Parallel(device):
    spec_layer = MFCC().to(device)
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    spec = spec_layer_parallel(x)    

@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])
def test_CQT1992_Parallel(device):
    spec_layer = CQT1992(fmin=110, n_bins=60, bins_per_octave=12).to(device)
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    spec = spec_layer_parallel(x) 

@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])    
def test_CQT1992v2_Parallel(device):
    spec_layer = CQT1992v2().to(device)
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    spec = spec_layer_parallel(x)        

@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])    
def test_CQT2010_Parallel(device):
    spec_layer = CQT2010().to(device)
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    spec = spec_layer_parallel(x)  
    
@pytest.mark.parametrize("device", [f'cuda:{gpu_idx}'])    
def test_CQT2010v2_Parallel(device):
    spec_layer = CQT2010v2().to(device)
    spec_layer_parallel = torch.nn.DataParallel(spec_layer)
    spec = spec_layer_parallel(x)