import torch import torch.nn as nn import numpy as np from time import time from ..librosa_functions import * from ..utils import * class VQT(torch.nn.Module): def __init__( self, sr=22050, hop_length=512, fmin=32.70, fmax=None, n_bins=84, filter_scale=1, bins_per_octave=12, norm=True, basis_norm=1, gamma=0, window='hann', pad_mode='reflect', earlydownsample=True, trainable=False, output_format='Magnitude', verbose=True ): super().__init__() self.norm = norm self.hop_length = hop_length self.pad_mode = pad_mode self.n_bins = n_bins self.earlydownsample = earlydownsample self.trainable = trainable self.output_format = output_format self.filter_scale = filter_scale self.bins_per_octave = bins_per_octave self.sr = sr self.gamma = gamma self.basis_norm = basis_norm # It will be used to calculate filter_cutoff and creating CQT kernels Q = float(filter_scale)/(2**(1/bins_per_octave)-1) # Creating lowpass filter and make it a torch tensor if verbose==True: print("Creating low pass filter ...", end='\r') start = time() lowpass_filter = torch.tensor(create_lowpass_filter( band_center = 0.50, kernelLength=256, transitionBandwidth=0.001) ) self.register_buffer('lowpass_filter', lowpass_filter[None,None,:]) if verbose == True: print("Low pass filter created, time used = {:.4f} seconds".format(time()-start)) n_filters = min(bins_per_octave, n_bins) self.n_filters = n_filters self.n_octaves = int(np.ceil(float(n_bins) / bins_per_octave)) if verbose == True: print("num_octave = ", self.n_octaves) self.fmin_t = fmin * 2 ** (self.n_octaves - 1) remainder = n_bins % bins_per_octave if remainder==0: # Calculate the top bin frequency fmax_t = self.fmin_t*2**((bins_per_octave-1)/bins_per_octave) else: # Calculate the top bin frequency fmax_t = self.fmin_t*2**((remainder-1)/bins_per_octave) # Adjusting the top minimum bins self.fmin_t = fmax_t / 2 ** (1 - 1 / bins_per_octave) if fmax_t > sr/2: raise ValueError('The top bin {}Hz has exceeded the Nyquist frequency, \ please reduce the n_bins'.format(fmax_t)) if self.earlydownsample == True: # Do early downsampling if this argument is True if verbose == True: print("Creating early downsampling filter ...", end='\r') start = time() sr, self.hop_length, self.downsample_factor, early_downsample_filter, \ self.earlydownsample = get_early_downsample_params(sr, hop_length, fmax_t, Q, self.n_octaves, verbose) self.register_buffer('early_downsample_filter', early_downsample_filter) if verbose==True: print("Early downsampling filter created, \ time used = {:.4f} seconds".format(time()-start)) else: self.downsample_factor = 1. # For normalization in the end # The freqs returned by create_cqt_kernels cannot be used # Since that returns only the top octave bins # We need the information for all freq bin alpha = 2.0 ** (1.0 / bins_per_octave) - 1.0 freqs = fmin * 2.0 ** (np.r_[0:n_bins] / np.double(bins_per_octave)) self.frequencies = freqs lenghts = np.ceil(Q * sr / (freqs + gamma / alpha)) # get max window length depending on gamma value max_len = int(max(lenghts)) self.n_fft = int(2 ** (np.ceil(np.log2(max_len)))) lenghts = torch.tensor(lenghts).float() self.register_buffer('lenghts', lenghts) my_sr = self.sr for i in range(self.n_octaves): if i > 0: my_sr /= 2 Q = float(self.filter_scale)/(2**(1/self.bins_per_octave)-1) basis, self.n_fft, lengths, _ = create_cqt_kernels(Q, my_sr, self.fmin_t * 2 ** -i, self.n_filters, self.bins_per_octave, norm=self.basis_norm, topbin_check=False, gamma=self.gamma) cqt_kernels_real = torch.tensor(basis.real.astype(np.float32)).unsqueeze(1) cqt_kernels_imag = torch.tensor(basis.imag.astype(np.float32)).unsqueeze(1) self.register_buffer("cqt_kernels_real_{}".format(i), cqt_kernels_real) self.register_buffer("cqt_kernels_imag_{}".format(i), cqt_kernels_imag) def forward(self, x, output_format=None, normalization_type='librosa'): """ Convert a batch of waveforms to VQT spectrograms. Parameters ---------- x : torch tensor Input signal should be in either of the following shapes.\n 1. ``(len_audio)``\n 2. ``(num_audio, len_audio)``\n 3. ``(num_audio, 1, len_audio)`` It will be automatically broadcast to the right shape """ output_format = output_format or self.output_format x = broadcast_dim(x) if self.earlydownsample==True: x = downsampling_by_n(x, self.early_downsample_filter, self.downsample_factor) hop = self.hop_length vqt = [] x_down = x # Preparing a new variable for downsampling my_sr = self.sr for i in range(self.n_octaves): if i > 0: x_down = downsampling_by_2(x_down, self.lowpass_filter) hop //= 2 else: x_down = x pad_length = int(getattr(self, 'cqt_kernels_real_{}'.format(i)).shape[-1] // 2) if self.pad_mode == 'constant': my_padding = nn.ConstantPad1d((pad_length, pad_length), 0) elif self.pad_mode == 'reflect': my_padding = nn.ReflectionPad1d((pad_length, pad_length)) cur_vqt = get_cqt_complex(x_down, getattr(self, 'cqt_kernels_real_{}'.format(i)), getattr(self, 'cqt_kernels_imag_{}'.format(i)), hop, my_padding) vqt.insert(0, cur_vqt) # print('vqt shape: ', vqt[0].shape) # PREVENT SIZE ALIGMENT ERROR # padding smaller frames, if the hop_length is not power of 2 # max_time_axis_frame = max(vqt_i.shape[-2] for vqt_i in vqt) min_time_axis_frame = min(vqt_i.shape[-2] for vqt_i in vqt) # print('max_time_axis_frame: ', max_time_axis_frame) # print('min_time_axis_frame: ', min_time_axis_frame) # vqt = [nn.functional.pad(vqt_i, (0, 0, max_time_axis_frame - vqt_i.shape[-2], 0), mode='constant', value=0) for vqt_i in vqt] vqt = [nn.functional.interpolate(vqt_i, size=(min_time_axis_frame, 2), mode='bilinear', align_corners=False) for vqt_i in vqt] # for vqt_i in vqt: # print(vqt_i.shape) vqt = torch.cat(vqt, dim=1) vqt = vqt[:,-self.n_bins:,:] # Removing unwanted bottom bins vqt = vqt * self.downsample_factor # Normalize again to get same result as librosa if normalization_type == 'librosa': vqt = vqt * torch.sqrt(self.lenghts.view(-1,1,1)) elif normalization_type == 'convolutional': pass elif normalization_type == 'wrap': vqt *= 2 else: raise ValueError("The normalization_type %r is not part of our current options." % normalization_type) if output_format=='Magnitude': if self.trainable==False: # Getting CQT Amplitude return torch.sqrt(vqt.pow(2).sum(-1)) else: return torch.sqrt(vqt.pow(2).sum(-1) + 1e-8) elif output_format=='Complex': return vqt elif output_format=='Phase': phase_real = torch.cos(torch.atan2(vqt[:,:,:,1], vqt[:,:,:,0])) phase_imag = torch.sin(torch.atan2(vqt[:,:,:,1], vqt[:,:,:,0])) return torch.stack((phase_real,phase_imag), -1)