Add main folder for repo

main/util.py

@@ -0,0 +1,317 @@
+from keras.callbacks import Callback
+import keras.backend as K
+import numpy as np
+
+
+class SGDRScheduler(Callback):
+    """custom callback for implementing a SGDR learning rate"""
+    def __init__(self, min_lr, max_lr, steps_per_epoch, lr_decay=0.9, cycle_length=10,
+                    mult_factor=1.5):
+        self.min_lr = min_lr
+        self.max_lr = max_lr
+        self.lr_decay = lr_decay
+
+        self.batch_since_restart = 0
+        self.next_restart = cycle_length
+
+        self.steps_per_epoch = steps_per_epoch
+
+        self.cycle_lenrfrrgth = cycle_length
+        self.mult_factor = mult_factor
+
+
+    def clr(self):
+        fraction_to_restart = self.batch_since_restart / (self.steps_per_epoch * self.cycle_length)
+        lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + np.cos(fraction_to_restart * np.pi))
+        return lr
+
+
+    def on_train_begin(self, logs=None):
+        K.set_value(self.model.optimizer.lr, self.max_lr)
+
+
+    def on_batch_end(self, batch, logs=None):
+        self.batch_since_restart += 1
+        K.set_value(self.model.optimizer.lr, self.clr())
+
+
+    def on_epoch_end(self, epoch, logs=None):
+        if epoch + 1 == self.next_restart:
+            self.batch_since_restart = 0
+            self.cycle_length = np.ceil(self.cycle_length * self.mult_factor)
+            self.next_restart += self.cycle_length
+            self.max_lr *= self.lr_decay
+
+
+# ----------------------------------------------------------------------------
+import matplotlib.pyplot as plt
+import keras.backend as K
+from keras.callbacks import Callback
+
+
+class LRFinder(Callback):
+    """
+    custom callback for evaluating the optimal lr range for SGDR
+    Usage:
+    lr_finder = models.util.LRFinder(min_lr=1e-5, max_lr=3e-2,
+        steps_per_epoch=np.ceil(n_sam/args.batch_size), epochs=3)
+    """
+    def __init__(self, min_lr=1e-5, max_lr=1e-2, steps_per_epoch=None, epochs=None):
+        super(LRFinder, self).__init__()
+
+        self.min_lr = min_lr
+        self.max_lr = max_lr
+        self.total_iterations = steps_per_epoch * epochs
+        self.iteration = 0
+        self.history = {}
+
+
+    def clr(self):
+        '''Calculate the learning rate.'''
+        x = self.iteration / self.total_iterations
+        return self.min_lr + (self.max_lr-self.min_lr) * x
+
+
+    def on_train_begin(self, logs=None):
+        '''Initialize the learning rate to the minimum value at the start of training.'''
+        logs = logs or {}
+        K.set_value(self.model.optimizer.lr, self.min_lr)
+
+
+    def on_batch_end(self, epoch, logs=None):
+        '''Record previous batch statistics and update the learning rate.'''
+        logs = logs or {}
+        self.iteration += 1
+
+        self.history.setdefault('lr', []).append(K.get_value(self.model.optimizer.lr))
+        self.history.setdefault('iterations', []).append(self.iteration)
+
+        for k, v in logs.items():
+            self.history.setdefault(k, []).append(v)
+
+        K.set_value(self.model.optimizer.lr, self.clr())
+
+
+    def plot_lr(self):
+        '''Helper function to quickly inspect the learning rate schedule.'''
+        plt.plot(self.history['iterations'], self.history['lr'])
+        plt.yscale('log')
+        plt.xlabel('Iteration')
+        plt.ylabel('Learning rate')
+        plt.tight_layout()
+        plt.savefig('plots/lr.png')
+        plt.clf()
+
+    def plot_loss(self):
+        '''Helper function to quickly observe the learning rate experiment results.'''
+        plt.plot(self.history['lr'], self.history['loss'])
+        plt.xscale('log')
+        plt.xlabel('Learning rate')
+        plt.ylabel('Loss')
+        plt.tight_layout()
+        plt.savefig('plots/loss.png')
+        plt.clf()
+
+
+# ----------------------------------------------------------------------------
+import tensorflow as tf
+import numpy as np
+import h5py
+import ds
+
+def load_data(args, type, num_files, full_data=False):
+    np.set_printoptions(threshold=100)
+    path = '../data/multispeaker'
+
+    # load training data
+    datasets = os.listdir(path)
+
+    for dataset in datasets:
+        if str(args.dim_size) and str(num_files) in dataset:
+            if args.new_data == 'False':
+                make_data = False
+                break
+        else:
+            make_data = True
+
+    if make_data:
+        ds.Prep_VCTK(type=type, num_files=num_files, dim=args.dim_size, file_list=f'{path}/{type}-files.txt')
+
+    with h5py.File(f'{path}/vctk-{type}.4.16000.{args.dim_size}.{num_files}.0.25.h5', 'r') as hf:
+        X = np.array(hf.get('data'))
+        Y = np.array(hf.get('label'))
+
+    n_sam, n_dim, n_chan = Y.shape
+    r = Y[0].shape[1] / X[0].shape[1]
+
+    if full_data:
+        return X, Y, n_sam
+    else:
+        return X, Y
+
+
+# ----------------------------------------------------------------------------
+import os
+
+
+def review_model(args, x_train, y_train):
+    """reviews model parameters and raises warnings if something is not recommended"""
+    # prints preivew of the data
+    preview_data(x_train, y_train)
+
+    # assert not overwriting weights
+    if args.from_ckpt == 'False':
+        files = os.listdir('./logs/weights')
+        for file in files:
+            if f'loss.{args.model_id}' in file:
+                input('Warning: Are you sure you want to write over these weights?')
+
+
+# ----------------------------------------------------------------------------
+import numpy as np
+
+
+def preview_data(X, Y):
+    print ('Preview X:')
+    print (f'Shape: {X.shape}')
+    print (f'Max: {np.amax(X)} | Min: {np.amin(X)}')
+    print (X[1])
+
+    print ('Preview Y:')
+    print (f'Shape of Y: {Y.shape}')
+    print (f'Max: {np.amax(Y)} | Min: {np.amin(Y)}')
+    print (Y[1])
+
+    # data = eval_wav.get_spectrum(X[:100].flatten(), n_fft=2048)
+    # label = eval_wav.get_spectrum(Y[:100].flatten(), n_fft=2048)
+
+    input('Press enter to continue...')
+
+
+# ----------------------------------------------------------------------------
+import os
+import librosa
+import numpy as np
+from keras.models import Model
+from scipy import interpolate
+from scipy.signal import decimate
+from matplotlib import pyplot as plt
+
+
+class eval_wav:
+    '''
+    Helper function for eval() in main.py
+    Takes a single wavfile and evaluates it by exporting audio and spectrogram
+        for hr, lr, and pr
+    '''
+    def __init__(self, file, args, model):
+        # ../data/VCTK-Corpus---
+        x_hr, fs = librosa.load(file, sr=args.sample_rate)
+
+        # ensure that input is a multiple of 2^downsampling layers
+        ds_layers = 5
+        x_hr = eval_wav.clip(x_hr, 2**ds_layers)
+        assert len(x_hr) % 2**ds_layers == 0
+
+        # downscale signal
+        # x_lr = decimate(x_hr, args.scale)
+        x_lr = np.array(x_hr[0::args.scale])
+        # x_lr = downsample_bt(x_hr, args.scale)
+        assert len(x_hr)/len(x_lr) == args.scale
+
+
+        # upsample signal through interpolation
+        x_ir = eval_wav.upsample(x_lr, args.scale)
+        assert len(x_ir) == len(x_hr)
+
+        # trim array again to make it a multiple of 800
+        x_ir = eval_wav.clip(x_ir, 800)
+        print(f'Input length: {len(x_ir)}')
+
+
+        n_sam = len(x_ir)/800
+        x_pr = model.predict(x_ir.reshape(int(n_sam), 800, 1))
+        x_pr = x_pr.flatten()
+
+        # save the file
+        filename = os.path.basename(file)
+        name = os.path.splitext(filename)[-2]
+
+        if args.make_audio:
+            audio_data = np.concatenate((x_hr, x_ir, x_pr), axis=0)
+            audio_outname = f'../samples/audio/{name}'
+            librosa.output.write_wav(audio_outname + '.hr.wav', audio_data, fs)
+
+        # save the spectrum
+        spec_outname = f'../samples/spectrograms/{name}'
+        self.outfile=spec_outname + '.png'
+        self.S_pr = eval_wav.get_spectrum(x_pr, n_fft=2048)
+        self.S_hr = eval_wav.get_spectrum(x_hr, n_fft=2048)
+        self.S_lr = eval_wav.get_spectrum(x_lr, n_fft=2048/args.scale)
+        self.S_ir = eval_wav.get_spectrum(x_ir, n_fft=2048)
+        self.save_spectrum()
+
+    @staticmethod
+    def upsample(x_lr, r): #lr = lowres, hr = highres
+
+      x_lr = x_lr.flatten() # flatten audio array
+      x_hr_len = len(x_lr) * r # get (len of audio array * scaling factor)
+      x_sp = np.zeros(x_hr_len) # create zero-padded array with new length
+
+      i_lr = np.arange(x_hr_len, step=r) # create lr array with step size of scaling factor
+      i_hr = np.arange(x_hr_len)
+
+      f = interpolate.splrep(i_lr, x_lr) # "Given the set of data points (x[i], y[i]) determine a smooth spline approximation"
+
+      # Given the knots and coefficients of a B-spline representation, evaluate the value of the smoothing polynomial and its derivatives.
+      x_sp = interpolate.splev(i_hr, f)
+
+      return x_sp
+
+    @staticmethod
+    def clip(array, multiple):
+        x_len = len(array)
+        remainder = x_len % multiple
+        x_len = x_len - remainder
+        array = array[:x_len]
+        return array
+
+    @staticmethod
+    def get_spectrum(data, n_fft=2048):
+        S = librosa.stft(data, int(n_fft))
+        S = np.log1p(np.abs(S))
+        p = np.angle(S)
+        S = np.log1p(np.abs(S))
+        return S.T
+
+    def save_spectrum(self, lim=1000):
+        plt.subplot(2,2,1)
+        plt.title('Target')
+        plt.xlabel('Frequency')
+        plt.ylabel('Time')
+        plt.imshow(self.S_hr, aspect=10)
+        plt.xlim([0,lim])
+
+        plt.subplot(2,2,2)
+        plt.title('Test')
+        plt.xlabel('Frequency')
+        plt.ylabel('Time')
+        plt.imshow(self.S_lr, aspect=10)
+        plt.xlim([0,lim])
+
+        plt.subplot(2,2,3)
+        plt.title('Interp')
+        plt.xlabel('Frequency')
+        plt.ylabel('Time')
+        plt.imshow(self.S_ir, aspect=10)
+        plt.xlim([0,lim])
+
+        plt.subplot(2,2,4)
+        plt.title('Predict')
+        plt.xlabel('Frequency')
+        plt.ylabel('Time')
+        plt.imshow(self.S_pr, aspect=10)
+        plt.xlim([0,lim])
+
+        plt.tight_layout()
+        plt.savefig(self.outfile)