VPM.py

import os
import csv

import scipy.io as sio
from scipy.io import wavfile
from scipy.io.wavfile import write
import scipy.signal as sis
import scipy.fftpack as fftpack

import librosa
import numpy as np
import librosa
import librosa.display
import matplotlib.pyplot as plt

from Utils import *

# Sample rate and bit depth
sample_rate = 44100
short_max = 32768

def create_data_ref_list(file_path, n_pitches, n_vowels, n_people):
    """Create a list of all the filenames in the dataset.

    To access the list of files, use data_ref_list[vowel_idx][pitch_idx]
    A specific filename is accessed with data_ref_list[vowel_idx][pitch_idx][person_idx]

    Args:
        filepath (str): The file path to the csv file containing the
            filenames of all .wav files
        n_pitches (int): The number of pitches in our dataset
        n_vowels (int): The number of vowels in our dataset
        n_people (int): The number of people in our dataset

    Returns:
        data_ref_list (list): A list of filenames, organized by word, then
            pitch, then people (n_vowels, n_pitches, n_people)
    """
    data_ref_list = [ [ [] for pIdx in range(0, n_pitches) ]
                      for wIdx in range(0, n_vowels) ]

    with open(file_path) as dataset_csv:
        reader = csv.reader(dataset_csv, delimiter=',')
        for idx, row in enumerate(reader):
            if idx == 0: continue
            filename, vowel_idx, pitch_idx, personNum = row
            data_ref_list[int(vowel_idx)][int(pitch_idx)].append(filename)
    return data_ref_list

def create_data_label_pairs(n_pitches):
    """Create a dictionary/array of data-label pairs.

    This provides an array of 3-tuples, as well as a dictionary of
    arrays, where each subarray contains 3-tuples, with elements:
    [shift_amt, input_pitch_idx, label_pitch_Idx],
    where input_pitch_idx is the input, label_pitch_Idx is the desired output.

    This is used so that we can feed data-label pairs to our neural net.
    We should use each 3-tuple n_people * n_vowels times, with
    the files referenced via data_list[vowel_idx][pitch_idx][person_idx],
    and the pitch shift amount provided to the NN.

    Args:
        n_pitches (int): The number of pitches in our dataset

    Returns:
        data_label_pairs: An array where each element is a 3-tuple of
            [shift_amt, input_pitch_idx, label_pitch_idx].
        data_label_pairs_dict: A dictionary of dimension n_pitchShifts, where each
            element is an array (n_words * n_startingPitches,), where:
            n_pitchShifts: the number of possible pitch shifts,
            n_startingPitches: the number of starting pitches for that
            shift_amt value
    """
    def append_tuple(shift_amt, pitch_idx):
        # [shift_amt, input_pitch_idx, label_pitch_Idx]
        data_label_pairs_dict[shift_amt].append(
            [shift_amt, pitch_idx, pitch_idx + shift_amt])
        data_label_pairs.append(
            [shift_amt, pitch_idx, pitch_idx + shift_amt])

    data_label_pairs = []
    data_label_pairs_dict = {}
    for pIdx in range(-n_pitches + 1, n_pitches):
        data_label_pairs_dict[pIdx] = []

    # Pitch indices range from 0-15, so we can shift from -15 to 15 pitches up.
    # First loop: shift_amt 0 to 15
    for shift_amt in range(0, n_pitches):
        # Iterate through available pitch shift starting points
        for pitch_idx in range(0, n_pitches - shift_amt):
            append_tuple(shift_amt, pitch_idx)
    # Second loop: shift_amt -15 to -1
    for shift_amt in range(-n_pitches + 1, 0):
        for pitch_idx in range(n_pitches - 1, -1 - shift_amt, -1):
            append_tuple(shift_amt, pitch_idx)

    return data_label_pairs, data_label_pairs_dict

def load_wav_files(rel_path, data_list):
    """Takes a list of filepaths, and returns a 2D array with all their data.

    The function also ensures that the resulting array only contains mono data.

    Args:
        rel_path (str): The relative path of the filenames in file_paths.
            i.e. the filepath would be rel_path/file_paths[i]
        file_paths (list): A list of filenames, where each filename is a
            .wav file to be added to the output.

    Returns:
        signal_data (np.ndarray): A 2D matrix that contains the audio data of
            all specified file paths, such that
            signal_data[i] provides the waveform of the ith file in file_paths.
            The dimensions are (len(file_paths), [length of file at filepath])
    """
    result = []
    for idx, file_path in enumerate(data_list):
        assert(file_path[-4:] == '.wav')
        s_r, short_data = sio.wavfile.read(os.path.join(rel_path,file_path))
        if (s_r != sample_rate):
            print("Please resample to 44100Hz for better results, but continuing regardless...");
        # Make it mono if it's stereo
        if len(short_data.shape) == 2 and short_data.shape[1] == 2:
            short_data = short_data[:, 0]
        result.append(short_data / short_max)
    return np.array(result)

def compute_hop_length(win_length, overlap):
    """Utility function to compute the hop_length.

    This is used by stft, istft, and ffts_to_melspectrogram.

    Args:
        win_length (int): The size of a window
        overlap (float): The amount of overlap between each window.

    Returns:
        hop_length (int): The computed hop_length.
    """
    return int(win_length * (1 - overlap))

def stft(waveform, win_length=1024, overlap=.5, window='hann', plot=True):
    """Takes a waveform and returns a 2D complex-valued matrix (spectrogram).

    The function performs STFT, i.e. windowing and performing FFT on each
    window. This is a wrapper for the librosa.core.stft function.

    Args:
        waveform (np.array): An array of amplitudes representing a signal.
        win_length (int): The size of each window (and corresponding FFT)
        overlap (float): The amount of overlap between each window. This
            translates to the hop_length.
        window (str): The window to use, specified by scipy.signal.get_window.
        plot (bool): If true, plot the spectrogram.

    Returns:
        ffts (np.ndarray): A 2D complex-valued matrix such that
            np.abs(ffts[f, t]) is the magnitude (of freq bin f at frame t)
            np.angle(ffts[f, t]) is the phase
            The dimensions are (win_length, [number of frames for waveform])
    """
    waveform_norm = librosa.util.normalize(waveform)
    hop_length = compute_hop_length(win_length, overlap)
    waveform_stft = librosa.core.stft(waveform_norm, n_fft=win_length, hop_length=hop_length, win_length=win_length, window=window)

    if plot:
        plot_ffts_spectrogram(waveform_stft, sample_rate)

    return waveform_stft

def istft(ffts, win_length=1024, overlap=.5, window='hann', save_file=False, file_name=''):
    """Takes a 2D complex-valued matrix (spectrogram) and returns a waveform.

    This function performs ISTFT, and is a wrapper for librosa.core.istft.

    Args:
        ffts (np.ndarray): A 2D complex-valued matrix such that ffts[f, t]
            is the complex number representing the FFT value for the spectrum
            of freq bin f at frame t.
            The dimensions are (win_length, [number of frames for waveform])
        win_length (int): The size of each window (and corresponding FFT)
        overlap (float): The amount of overlap between each window. This
            translates to the hop_length.
        window (str): The window to use, specified by scipy.signal.get_window
        save_file (bool): If true, save of waveform to audio wav file.
        file_name (str): The respective file name for the audio wav file.

    Returns:
        waveform (np.array): An array of amplitudes representing a signal.
    """
    hop_length = compute_hop_length(win_length, overlap)
    waveform_istft = librosa.core.istft(ffts, hop_length=hop_length, win_length=win_length, window=window)

    if save_file:
        # change the path and file name accordingly
        librosa.output.write_wav('output_wav/' + file_name + '.wav', waveform_istft, sample_rate)

    return waveform_istft

def ffts_to_mel(ffts, n_mels=128, n_mfcc=20, skip_mfcc=False, plot=False):
    """Converts a spectrogram to a mel-spectrogram and MFCC.

    This function is a wrapper for librosa.feature.melspectrogram and
    librosa.feature.mfcc.

    Args:
        ffts (np.ndarray): A 2D complex-valued matrix such that ffts[f, t]
            is the complex number representing the FFT value for the spectrum
            of freq bin f at frame t.
            The dimensions are (win_length, [number of frames for waveform])
        n_mels (int): The number of Mel bands to generate.
        n_mfcc (int): The number of MFCC features to compute.
        skip_mfcc(boolean): If True, do compute the MFCC, as we only want the
            spectrogram information.

    Returns:
        mel_spectrogram (np.ndarray): A 2D matrix such that
            mel_spectrogram[m, t] is the magnitude of mel bin m at frame t
            The dimensions are (n_mels, [ffts.shape[1]])
        mfcc (np.ndarray): A 2D matrix such that
            mfcc[m, t] is the magnitude of the mth feature at frame t
    """
    """
    !! Write code here !!
    Louiz's note: Please handle sampling rate properly, we assume always 44100.
    Check out librosa.filters.mel if unsure how to write the arguments to call
    librosa.feature.melspectrogram.
    """
    D = np.abs(ffts) ** 2
    mel_freq_spec = librosa.feature.melspectrogram(S=D, sr=sample_rate, n_mels=n_mels)

    if plot:
        plot_mel_spectrogram(mel_freq_spec, sample_rate)

    if not skip_mfcc:
        mfccs = librosa.feature.mfcc(S=librosa.power_to_db(mel_freq_spec), sr=sample_rate, n_mfcc=n_mfcc)

        if plot:
            plot_mfcc(mfccs, sample_rate)

        return mel_freq_spec, mfccs

    return mel_freq_spec

def simple_fft_pitch_shift(fft, shift_amt, n_ffts=1024):
    """Takes a single fft vector and shifts all values in the frequency domain.

    This is a "naive" pitch shift that simply up-shifts the values in the
    given fft, and is not expected to sound natural. Note that this is done on
    a SINGLE time slice. For shifting of an entire spectrogram, use
    simple_ffts_pitch_shift instead.

    Explanation:
    Assume that each value in fft (e.g. fft[f]), is given by a (value, freq)
    pair. We shift the pitch by multiplying the frequency values by
    (2**(shift_amt/12)), and interpolating the values back into the original
    frequency bin values (since the fft bins must keep their original frequency
    resolution).

    Example:
    Assume a frequency resolution of 20Hz, where bin 0: 0Hz, bin 1: 20 Hz etc.
    If z = fft[3], where bin 3 is the frequency 60Hz, we denote this
    as (z, 60). We also have z' = fft[2], denoted as (z', 40).

    So shifting the pitch by 1, we end up with (z, 63.57) and (z', 42.4).
    So to get the value at 60Hz, we will need to interpolate between z' to z.

    Args:
        fft (np.array): A complex-valued array such that ffts[f] is the
            complex number representing the FFT value for the spectrum of freq
            bin f. Dimensions are (win_length,), where win_length is the
            number of windows used to generate this fft.
        shift_amt (int): The number of semitones to shift the pitch by. The
            expected range is [-15, 15].
        n_ffts (int): The number of fft bins.

    Returns:
        shifted_fft (np.array): A (win_length,) array with the shifted fft.
    """
    assert(-15 <= shift_amt and shift_amt <= 15)

    freqs = librosa.core.fft_frequencies(sample_rate,n_ffts)
    shifted_freqs = freqs * np.power(2, shift_amt/12)
    shifted_fft = np.interp(freqs, shifted_freqs, fft)

    return shifted_fft

def simple_ffts_pitch_shift(ffts, shift_amt, n_ffts=1024):
    """Takes a 2D spectrogram, and pitch_shifts each time slice.

    Args:
        ffts (np.ndarray): A 2D complex-valued matrix such that ffts[f, t]
            is the complex number representing the FFT value for the spectrum
            of freq bin f at frame t.
            The dimensions are (win_length, [number of frames for waveform])
        shift_amt (int): The number of semitones to shift the pitch by. The
            expected range is [-15, 15].
        n_ffts (int): The number of fft bins.
    Returns:
        shifted_ffts (np.ndarray): A spectrogram of equal dimensions to ffts,
            with shifted frequency space.
    """
    assert(-15 <= shift_amt and shift_amt <= 15)
    return np.array([ simple_fft_pitch_shift(fft, shift_amt, n_ffts) for fft in ffts.T ]).T

def compute_new_sample_rate(base_sample_rate, shift_amt):
    """Returns a new sample rate based on a pitch shift amount.
    This is used by resample_wavs.
    Args:
        base_sample_rate (int): The original sampling rate
        shift_amt (int): The number of semitones to pitch shift by.
    Returns:
        new_sample_rate (int): The new sampling rate
        factor (float): The percentage time compression to achieve the increase in pitch
            specified by shift_amt
    """
    factor = 2 ** (-shift_amt/12)
    new_sample_rate = base_sample_rate * factor

    return new_sample_rate, factor

def resample_wavs(all_wav, shift_amt):
    """Resamples wav files, which results in pitch shift + time compression/stretch.
    Args:
        all_wav (np.ndarray): A 2D matrix such that all_wav[i] provides the ith waveform
        shift_amt (int): The number of semitones that all wav files should be shifted by.
    Returns:
        shifted_wavs (np.ndarray): 2D matrix where each wav in all_wav is resampled.
    """
    new_sample_rate, factor = compute_new_sample_rate(sample_rate, shift_amt)
    shifted_wavs = np.array([librosa.resample(track_wav, sample_rate, new_sample_rate) for track_wav in all_wav])
    return shifted_wavs

def stretch_wavs(all_wav, shift_amt, overlap, n_ffts=1024):
    """Stretches wav files by shift_amt. Takes a 2D matrix.
    Args:
        all_wav (np.ndarray): A 2D matrix such that all_wav[i] provides the ith waveform
        shift_amt (int): The number of semitones that all wav files should be shifted by.
    Returns:
        all_wavs_stretched (np.ndarray): 2D matrix where each wav in all_wav is stretched.
    """
    all_ffts = np.array([ stft(waveform, win_length=n_ffts, overlap=overlap, plot=False) 
                          for waveform in all_wav ])
    new_sample_rate, factor = compute_new_sample_rate(sample_rate, int(shift_amt))
    stretched_ffts = np.array([ librosa.phase_vocoder(ffts, factor, hop_length=compute_hop_length(n_ffts, overlap)) 
                                for ffts in all_ffts ])
    all_wavs_stretched = np.array( [ istft(ffts, overlap=overlap, win_length=n_ffts, save_file=False) 
                                   for ffts in stretched_ffts ])
    return all_wavs_stretched

def resample_pitch_shift(all_wav, shift_amt, overlap, n_ffts=1024):
    """Pitch shifts wav files via resampling. Returns both shifted wav and fft data.
    
    First, phase_vocoder is used to stretch/compress the wav files, then they are resampled
    to result in a pitch shift without change in length.

    Args:
        all_wav (np.ndarray): A 2D matrix such that all_wav[i] provides the ith waveform
        shift_amt (int): The number of semitones that all wav files should be shifted by.
        overlap (float): The amount of overlap between each window.
        n_ffts (int): The number of FFT bins.
    Returns: 
        pitch_shifted_data (np.ndarray): 2D matrix where each wav in all_wav is pitched up.
        pitched_spectra (np.ndarray): 3D matrix where each entry is the STFT of each pitch
            shifted wav file.
    """
    all_stretched_data = stretch_wavs(all_wav, shift_amt, overlap, n_ffts=n_ffts)
    pitch_shifted_data = resample_wavs(all_stretched_data, shift_amt)
    pitched_spectra = \
        (np.array([ stft(waveform, win_length=n_ffts, overlap=overlap, plot=False) 
                    for waveform in pitch_shifted_data ]))
    return pitch_shifted_data, pitched_spectra