<reponame>gundersena/audio-enhancer
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)

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)
custom callback for implementing a SGDR learning rate # ---------------------------------------------------------------------------- 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) Calculate the learning rate. Initialize the learning rate to the minimum value at the start of training. Record previous batch statistics and update the learning rate. Helper function to quickly inspect the learning rate schedule. Helper function to quickly observe the learning rate experiment results. # ---------------------------------------------------------------------------- # load training data # ---------------------------------------------------------------------------- reviews model parameters and raises warnings if something is not recommended # prints preivew of the data # assert not overwriting weights # ---------------------------------------------------------------------------- # data = eval_wav.get_spectrum(X[:100].flatten(), n_fft=2048) # label = eval_wav.get_spectrum(Y[:100].flatten(), n_fft=2048) # ---------------------------------------------------------------------------- Helper function for eval() in main.py     Takes a single wavfile and evaluates it by exporting audio and spectrogram         for hr, lr, and pr # ../data/VCTK-Corpus--- # ensure that input is a multiple of 2^downsampling layers # downscale signal # x_lr = decimate(x_hr, args.scale) # x_lr = downsample_bt(x_hr, args.scale) # upsample signal through interpolation # trim array again to make it a multiple of 800 # save the file # save the spectrum #lr = lowres, hr = highres # flatten audio array # get (len of audio array * scaling factor) # create zero-padded array with new length # create lr array with step size of scaling factor # "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.