<filename>modules/audio/transforms.py

import torch
import torch.nn as nn
import torch.nn.functional as F

# stft module with istft support
class STFT (nn.Module):

    def __init__ (self,
        n_fft=400,
        hop_length=None,
        win_length=None):

        super(STFT,self).__init__()

        self.n_fft=n_fft,
        self.hop_length=hop_length,
        self.win_length=win_length,
        self.window=torch.hann_window(win_length),

    def forward (self,sample):

        return torch.stft(sample,self.n_fft,self.hop_length,self.win_length,self.window,True,'reflect',False,True)

    # torch.istft not implemented yet
    # this is a barebones version of [https://gist.github.com/keunwoochoi/2f349e72cc941f6f10d4adf9b0d3f37e]
    def inverse (self,stft):

        sample = torch.zeros(
            (stft.shape[0],self.n_fft + self.hop_length * (stft.shape[-2] - 1)),
            device=stft.device)

        for i in range(stft.shape[-2]):

            partial = torch.irfft(stft[:,:,i],signal_ndim=1,signal_sizes=(self.win_length))
            partial = self.window * partial

            sample[:,i * self.hop_length:i * self.hop_length + self.n_fft] += partial

        sample = sample[:,self.n_fft // 2:]
        sample = sample / self.n_fft / self.hop_length / 2.0

        return sample

# mel scaling module
class MelScale (nn.Module):

    def __init__ (self,
        n_mels=128,
        sample_rate=16000,
        f_min=0,
        f_max=None,
        n_stft=400 // 2 + 1):

        super(MelScale,self).__init__()

        self.conversion_matrix = self.conversion_matrix(n_stft,f_min,f_max,n_mels)

    def forward (self,sample):

        return torch.matmul(sample,self.conversion_matrix)

    def scale_hertz_to_mel (self,hertz):

        return 2595 * torch.log10(1.0 + (hertz / 700))

    def scale_mel_to_hertz (self,mel):

        return 700 * (10 ** (mel / 2595) - 1.0)

    def conversion_matrix (self,n_stft,f_min,f_max,n_mels):

        frequencies = torch.linspace(f_min,f_max,n_stft)

        mel_min = self.scale_hertz_to_mel(f_min)
        mel_max = self.scale_hertz_to_mel(f_max)

        frequencies_points = torch.linspace(m_min,m_max,n_mels + 2)
        frequencies_points = self.scale_mel_to_hertz(frequencies_points)

        frequency_difference = frequencies_points[1:] - frequencies_points[:-1]
        slopes = frequencies_points.unsqueeze(0) - frequencies.unsqueeze(1)

        zero = torch.tensor(0,dtype=slopes.dtype,device=slopes.device)
        up_slopes = slopes[:,2:] / frequency_difference[1:]
        down_slopes = (-1 * slopes[:,:-2]) / frequency_difference[:-1]

        return torch.max(zero,torch.min(up_slopes,down_slopes))

# instantaneous frequency module [https://magenta.tensorflow.org/gansynth]
class InstantaneousFrequency (nn.Module):

    def __init__ (self):

        super(InstantaneousFrequency,self).__init__()

    # torch port of numpy.diff
    def difference (self,phase):

        return phase[:,1:,:] - phase[:,:-1,:]

    # torch port of numpy.unwrap
    def unwrap (self,phase):

        phase_difference = self.difference(phase)

        phase_difference_modulo = torch.remainder(phase_difference + np.pi, 2.0 * np.pi) - np.pi
        phase_difference_modulo[(phase_difference_modulo == -np.pi) & (phase_difference > 0)] = np.pi

        phase_correct = phase_difference_modulo - phase_difference
        phase_correct[torch.abs(phase_difference) < np.pi] = 0

        phase[:,1:] += torch.cumsum(phase_correct,dim=1)

        return phase

    # torch port of instantaneous frequency [https://github.com/tensorflow/magenta/blob/cc18fe9feba5fecc56a540864b15c88d33aaa028/magenta/models/gansynth/lib/spectral_ops.py#L199]
    def forward (self,phase):

        phase = self.unwrap(phase)
        phase[:,1:,:] = self.difference(phase) / np.pi

        return phase
