u-vector.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "u-vector.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y2YFbKGl3IjI"
      },
      "source": [
        "import os, librosa, time, pickle, random, warnings\n",
        "from glob import glob\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "tf.get_logger().setLevel('ERROR')\n",
        "warnings.filterwarnings('ignore')\n",
        "import tensorflow_io as tfio\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.ticker import MaxNLocator\n",
        "import seaborn as sns\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "import IPython.display as ipd\n",
        "from IPython.core.display import display, clear_output\n",
        "#%load_ext tensorboard\n",
        "\n",
        "from sklearn.manifold import TSNE\n",
        "from sklearn.preprocessing import OneHotEncoder, LabelEncoder\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "from tensorflow.keras.layers import *\n",
        "from tensorflow.keras.applications import *\n",
        "from tensorflow.keras import *\n",
        "\n",
        "from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score\n",
        "from sklearn.cluster import KMeans\n",
        "from coclust.evaluation.external import accuracy\n",
        "\n",
        "SR = 16_000\n",
        "FRAME = 0.2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MVeuawe53PjC"
      },
      "source": [
        "def loadData(name, frame=0.2, seconds=None):\n",
        "    '''\n",
        "    Dataset name prefixes: \n",
        "    TIMIT : DARPA TIMIT\n",
        "    LIBRI : LibriSpeech\n",
        "    ASR : Bengali ASR\n",
        "    '''\n",
        "\n",
        "    basepath = \"./drive/My Drive/SpeakerRecognition/preprocessed_dataset\"\n",
        "\n",
        "    data_types = [f\"{name}_X_16000_{seconds}_{frame}.pkl\", \n",
        "                  f\"{name}_y_16000_{seconds}_{frame}.pkl\",\n",
        "                  f\"{name}_fy_16000_{seconds}_{frame}.pkl\", \n",
        "                  f\"noise_16000_{frame}.pkl\"]\n",
        "\n",
        "    rets = []\n",
        "    for i in range(len(data_types)):\n",
        "        with open(os.path.join(basepath, data_types[i]), 'rb') as f:\n",
        "            rets.append(pickle.load(f))\n",
        "\n",
        "    X = np.asarray(rets[0], dtype=np.float32)\n",
        "    X = np.expand_dims(X, axis=-1)\n",
        "\n",
        "    y = LabelEncoder().fit_transform(rets[1])\n",
        "    fi = LabelEncoder().fit_transform(rets[2])\n",
        "\n",
        "    noise = np.asarray(rets[3], dtype=np.float32)\n",
        "    noise = np.expand_dims(noise, axis=-1)\n",
        "\n",
        "    # Returns X, y, file_indexes, noise \n",
        "    print(\"Total unique labels:\", len(np.unique(y)))\n",
        "    return X, y, fi, noise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F_NYnU-e3V7l"
      },
      "source": [
        "def speakerFilter(data_x, data_y, speakers, data_fi=None, seed=42):\n",
        "    '''\n",
        "    Limits number of speakers\n",
        "    data_x  : Speech frame\n",
        "    data_y  : actual label of speech frames\n",
        "    speakers: number of speakers\n",
        "    '''\n",
        "\n",
        "    tot_spkr = len(np.unique(data_y))\n",
        "    rand = random.Random(seed)\n",
        "    persons = rand.sample(range(tot_spkr), speakers)\n",
        "    #print(persons)\n",
        "\n",
        "    idx = np.asarray([i for i in range(data_x.shape[0]) if data_y[i] in persons], \n",
        "                     dtype=np.int32)\n",
        "\n",
        "    if data_fi is not None:\n",
        "        return (data_x[idx], LabelEncoder().fit_transform(data_y[idx]), \n",
        "                data_fi[idx])\n",
        "    return (data_x[idx], LabelEncoder().fit_transform(data_y[idx]))\n",
        "\n",
        "\n",
        "# An utility function\n",
        "def showWAV(wav, sr, frame):\n",
        "    '''\n",
        "       showWAV(X[:2], sr=SR, frame=FRAME)\n",
        "       wav shape : (x, y, z)\n",
        "    '''\n",
        "    plt.figure(figsize=(12, 3))\n",
        "    for i in range(wav.shape[0]):\n",
        "        plt.subplot(wav.shape[0], 1, i+1)\n",
        "        display(ipd.Audio(wav[i, ..., 0], rate=SR, autoplay=False))\n",
        "        plt.plot(np.arange(int(sr*frame)), wav[i, ..., 0])\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def makeImpure(data_y, ratio=0, seed=42):\n",
        "    if ratio == 0:\n",
        "        return data_y\n",
        "\n",
        "    tot_spkr = len(np.unique(data_y))\n",
        "    tot_impurs = int(data_y.shape[0]*ratio)\n",
        "    rand = random.Random(seed)\n",
        "\n",
        "    idxs = rand.sample(range(data_y.shape[0]), tot_impurs)\n",
        "    shuffle_data = data_y[idxs]\n",
        "    rand.shuffle(shuffle_data)\n",
        "\n",
        "    ret_data = np.copy(data_y)\n",
        "    ret_data[idxs] = shuffle_data\n",
        "\n",
        "    return ret_data\n",
        "\n",
        "\n",
        "def pairwiseRelations(data_y, max_lim, seed=42):\n",
        "    '''\n",
        "    Limits & constructs the number of pairwise relations\n",
        "    data_y  : The actual label of speech frames\n",
        "    max_lim : Maximum number of pairwise relations\n",
        "    '''\n",
        "\n",
        "    tot_spkr = len(np.unique(data_y))\n",
        "    rand = random.Random(seed)\n",
        "\n",
        "    av_idx = dict()\n",
        "    for i in range(data_y.shape[0]):\n",
        "        label = data_y[i]\n",
        "        if label not in av_idx:\n",
        "            av_idx[label] = set()\n",
        "        av_idx[label].add(i)\n",
        "    \n",
        "    tot_buckets = data_y.shape[0]//max_lim\n",
        "    bucket = [0 for i in range(tot_buckets)]\n",
        "    label = 0\n",
        "\n",
        "    new_y = np.zeros(data_y.shape, dtype=np.int32)\n",
        "\n",
        "    for i in range(tot_buckets):\n",
        "        for l in range(tot_spkr):\n",
        "            while len(av_idx[l]) > 0:\n",
        "                min_lim = len(av_idx[l])\n",
        "                ids = rand.sample(av_idx[l], min(min_lim, max_lim))\n",
        "                new_y[ids] = label\n",
        "                label += 1\n",
        "                for id in ids: av_idx[l].remove(id)\n",
        "\n",
        "    return new_y"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aImG4BKq3WmQ"
      },
      "source": [
        "# FFT was not used!!\n",
        "class FFT(tf.keras.layers.Layer):\n",
        "    def __init(self):\n",
        "        super(FFT, self).__init__()\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        super(FFT, self).build(input_shape)\n",
        "\n",
        "    def adapt(self, input):\n",
        "        return self.call(input)\n",
        "\n",
        "    def call(self, input):\n",
        "        # Since tf.signal.fft applies FFT on the innermost dimension,\n",
        "        # we need to squeeze the dimensions and then expand them again\n",
        "        # after FFT\n",
        "        input = tf.squeeze(input, axis=-1)\n",
        "        fft = tf.signal.fft(\n",
        "            tf.cast(tf.complex(real=input, imag=tf.zeros_like(input)), tf.complex64)\n",
        "        )\n",
        "        fft = tf.expand_dims(fft, axis=-1)\n",
        "        # Return the absolute value of the first half of the FFT\n",
        "        # which represents the positive frequencies\n",
        "        return tf.math.abs(fft[:, : (input.shape[1] // 2), :])\n",
        "\n",
        "\n",
        "class MEL(tf.keras.layers.Layer):\n",
        "    def __init__(self, scale=None, setup=2, custom_setup=None, \n",
        "                 output_channel=1):\n",
        "        '''\n",
        "            scale :        Can be [None, 'spec', 'log', 'db']\n",
        "                           By default it returns mel\n",
        "            custom_setup : Select parameter setups. must be in dict format\n",
        "            setup :        Select parameter setups that are predefined. \n",
        "                           Default is setup=2\n",
        "            output_channel: Number of output channels. If greater than one, the\n",
        "                            output will be repeated.\n",
        "        '''\n",
        "        self.scale = scale\n",
        "        super(MEL, self).__init__()\n",
        "        self.setup = [{\"nfft\": 191, \"window\": 128, \"stride\": 34, \"mels\": 100},\n",
        "                      {\"nfft\": 1024, \"window\": 128, \"stride\": 61, \"mels\": 263},\n",
        "                      {\"nfft\": 511, \"window\": 32, \"stride\": 16, \"mels\": 256}]\n",
        "        self.id = setup\n",
        "        self.output_channel = output_channel \n",
        "        if custom_setup is not None:\n",
        "            self.id = 0\n",
        "            self.setup = [custom_setup]\n",
        "        \n",
        "    def build(self, input_shape):\n",
        "        super(MEL, self).build(input_shape)\n",
        "\n",
        "    def adapt(self, input):\n",
        "        return self.call(input)\n",
        "\n",
        "    def call(self, input):\n",
        "        spect = tfio.experimental.audio.spectrogram(input[..., 0], \n",
        "                                                    nfft=self.setup[self.id][\"nfft\"], \n",
        "                                                    window=self.setup[self.id][\"window\"], \n",
        "                                                    stride=self.setup[self.id][\"stride\"])\n",
        "        \n",
        "        if self.scale == \"spec\":\n",
        "            spect = tf.expand_dims(spect, axis=-1)\n",
        "            if self.output_channel > 1:\n",
        "                return tf.keras.backend.repeat_elements(spect, \n",
        "                                                        self.output_channel, \n",
        "                                                        axis=-1)\n",
        "            return spect\n",
        "\n",
        "        mel = tfio.experimental.audio.melscale(spect, rate=SR, \n",
        "                                               mels=self.setup[self.id][\"mels\"], \n",
        "                                               fmin=0, fmax=8000)\n",
        "\n",
        "        if self.scale == \"log\":\n",
        "            mel = tf.math.log(mel)\n",
        "        elif self.scale == \"db\":\n",
        "            mel = tfio.experimental.audio.dbscale(mel, top_db=128)\n",
        "\n",
        "        mel = tf.expand_dims(mel, axis=-1)\n",
        "        if self.output_channel > 1:\n",
        "            mel = tf.keras.backend.repeat_elements(mel, self.output_channel, \n",
        "                                                    axis=-1)\n",
        "        return mel"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5Hvg-UJ23ZpW"
      },
      "source": [
        "class MyLogger(tf.keras.callbacks.Callback):\n",
        "    '''\n",
        "    Parameters:\n",
        "\n",
        "    n    : Number of steps after which logs will appear/calculated\n",
        "    plot : If True, the graphs (ACC graph, scatters) will be plotted \n",
        "    scatter : If True, the scatter will be plotted, plot must be True\n",
        "    val_train : The metrices will be also be calculated for the trainind data and\n",
        "              the corresponding pseudo labels\n",
        "    AE : The AutoEmbedder portion of the model, used for evaluation\n",
        "    \n",
        "    validation_data : [(X, pseudo_labels), (X, actual_labels)]\n",
        "    nodes : Number of data nodes that will be used for evaluation\n",
        "            default is 500\n",
        "\n",
        "    save_model : If True, the model will be saved when max ACC on actual label\n",
        "                 is found\n",
        "    \n",
        "    savepath : the basepath of the save directory\n",
        "\n",
        "    * self.start_epoch contains the last epoch when the training is terminated\n",
        "    '''\n",
        "    def __init__(self, validation_data, n=1, plot=False, save_model=True,\n",
        "                 savepath=None, dg=None, AE=None, val_train=False, \n",
        "                 scatter=False, nodes=500, show_fig=True,\n",
        "                 # MEL layer setup\n",
        "                 mel_scale='spec', setup=0, output_channel=1,):\n",
        "        \n",
        "        self.n = n\n",
        "        if validation_data != None:\n",
        "            self.x_tr, self.y_tr = validation_data[0]\n",
        "            self.x_val, self.y_val = validation_data[1]\n",
        "            self.classes_val = len(np.unique(self.y_val))\n",
        "            self.classes_tr = len(np.unique(self.y_tr))\n",
        "        self.start_time = time.time()\n",
        "        self.savepath = savepath\n",
        "        self.plot = plot\n",
        "        self.save_model = save_model\n",
        "        self.maxACC = 0\n",
        "        self.start_epoch = 0\n",
        "        self.nodes = nodes\n",
        "        self.AE = AE\n",
        "        self.val_train = val_train\n",
        "        self.scatter = scatter\n",
        "        self.show_fig = show_fig\n",
        "        self.mel = MEL(scale=mel_scale, setup=setup, \n",
        "                       output_channel=output_channel)\n",
        "\n",
        "        self.savelog = {'Epoch': [], 'ACC':[], 'NMI':[], 'ARI':[],\n",
        "                        'val_Epoch': [], 'val_ACC':[], 'val_NMI':[], \n",
        "                        'val_ARI':[], 'loss':[]}\n",
        "        \n",
        "        # Defining save paths\n",
        "        if self.savepath != None:\n",
        "            self.logpath = os.path.join(self.savepath, 'log.pickle')\n",
        "            self.modelpath = os.path.join(self.savepath, 'model', '')\n",
        "        # Creating save paths\n",
        "        if self.savepath != None and os.path.exists(self.savepath) == False:\n",
        "            os.makedirs(self.savepath)\n",
        "        # Loading previous data if found\n",
        "        if self.savepath != None and os.path.exists(self.logpath):\n",
        "            with open(self.logpath, 'rb') as f:\n",
        "                self.savelog = pickle.load(f)\n",
        "            print('Previous data loaded, starting epoch:', self.savelog['Epoch'][-1])\n",
        "            self.start_epoch = self.savelog['Epoch'][-1]\n",
        "            self.maxACC = max(self.savelog['val_ACC'])\n",
        "        if self.plot:\n",
        "            sns.set_style(\"whitegrid\")\n",
        "\n",
        " \n",
        "    def on_epoch_end(self, epoch, logs={}):\n",
        "        if epoch % self.n != 0: return\n",
        "        epoch += 1\n",
        "\n",
        "        # Validation data\n",
        "        vacc, vnmi, vari, vouts, vys = self._KmeansAcc(self.x_val, self.y_val, \n",
        "                                           self.classes_val)\n",
        "        logs['val_acc'], logs['val_nmi'], logs['val_ari'] = vacc, vnmi, vari\n",
        "\n",
        "        # Train data\n",
        "        if self.val_train:\n",
        "            acc, nmi, ari, outs, ys = self._KmeansAcc(self.x_tr, self.y_tr, \n",
        "                                            self.classes_tr)\n",
        "            logs['acc'], logs['nmi'], logs['ari'] = acc, nmi, ari\n",
        "        else:\n",
        "            acc, nmi, ari = None, None, None\n",
        "\n",
        "        ep_time = time.time() - self.start_time\n",
        "        self._saveLog(epoch, acc, nmi, ari, logs['loss'], vacc, vnmi, vari)\n",
        "        if self.plot: self.plotter(vouts, vys)\n",
        "        else : clear_output(wait=True)\n",
        "\n",
        "        self.start_time = time.time()\n",
        "        print(f\"Epoch {epoch}: bACC: {self.maxACC:.2f}\")\n",
        "        print(f\"val_ACC {vacc:.3f} val_NMI {vnmi:.3f} val_ARI {vari:.3f} Loss {logs['loss']:.3f} ET:{ep_time:.1f}\")\n",
        "        if self.val_train:\n",
        "            print(f\"tACC {acc:.3f} tNMI {nmi:.3f} tARI {ari:.3f}\")\n",
        "\n",
        "        if self.savepath != None and vacc > self.maxACC:\n",
        "            print(f'Saving model, prev :{self.maxACC:.2f}, current: {vacc:.2f}')\n",
        "            self.maxACC = vacc\n",
        "            with open(self.logpath, 'wb') as f:\n",
        "                pickle.dump(self.savelog, f)\n",
        "            if self.save_model:\n",
        "                tf.keras.models.save_model(self.model, self.modelpath)\n",
        "        elif self.savepath != None:\n",
        "            with open(self.logpath, 'wb') as f:\n",
        "                pickle.dump(self.savelog, f)\n",
        "                \n",
        "        self.maxACC = max(vacc, self.maxACC)\n",
        "\n",
        "\n",
        "    def _saveLog(self, epoch, acc, nmi, ari, loss, vacc, vnmi, vari):        \n",
        "        self.savelog['Epoch'].append(epoch)\n",
        "        # Train data\n",
        "        if self.val_train:\n",
        "            self.savelog['ACC'].append(acc)\n",
        "            self.savelog['NMI'].append(nmi)\n",
        "            self.savelog['ARI'].append(ari)\n",
        "        # Validation data\n",
        "        self.savelog['val_ACC'].append(vacc)\n",
        "        self.savelog['val_NMI'].append(vnmi)\n",
        "        self.savelog['val_ARI'].append(vari)\n",
        "        self.savelog['loss'].append(loss)\n",
        " \n",
        " \n",
        "    def plotter(self, embds, labels):\n",
        "        clear_output(wait=True)\n",
        "        plt.figure(figsize=(12.8, 4.8))\n",
        "        if self.scatter:\n",
        "            plt.subplot(1, 2, 1)\n",
        "        \n",
        "        plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))\n",
        "        # dashes are the train data\n",
        "        if self.val_train:\n",
        "            plt.plot(self.savelog['Epoch'], (self.savelog['ACC']),\n",
        "                    '--', label='tACC', c='purple')\n",
        "            plt.plot(self.savelog['Epoch'], (self.savelog['NMI']), \n",
        "                    '--', label='tNMI', c='orange')\n",
        "            plt.plot(self.savelog['Epoch'], (self.savelog['ARI']), \n",
        "                    '--', label='tARI', c='c')\n",
        "        plt.plot(self.savelog['Epoch'], \n",
        "                 (np.array(self.savelog['loss'])/max(self.savelog['loss'])), \n",
        "                 label='Training Loss', c='r')\n",
        "        # solids are validation data\n",
        "        plt.plot(self.savelog['Epoch'], (self.savelog['val_ACC']), \n",
        "                 label='ACC', c='purple')\n",
        "        plt.plot(self.savelog['Epoch'], (self.savelog['val_NMI']), \n",
        "                 label='NMI', c='orange')\n",
        "        plt.plot(self.savelog['Epoch'], (self.savelog['val_ARI']), \n",
        "                 label='ARI', c='c')\n",
        "        plt.ylabel('Score')\n",
        "        plt.xlabel('Epoch')\n",
        "\n",
        "        plt.grid()\n",
        "        plt.minorticks_on()\n",
        "        plt.grid(b=True, which='minor', linestyle='--', alpha=0.25)\n",
        "        plt.legend(loc='upper center', bbox_to_anchor=(0.5, -0.13),\n",
        "                   fancybox=True, shadow=False, ncol=5)\n",
        "\n",
        "        if self.scatter:\n",
        "            plt.subplot(1, 2, 2)\n",
        "            p = TSNE(n_components=2).fit_transform(embds)\n",
        "            plt.scatter(p[:, 0], p[:, 1], c=labels)\n",
        " \n",
        "        if self.savepath != None:\n",
        "            plt.savefig(os.path.join(self.savepath, \"logPlot.png\"), \n",
        "                        transparent=True, bbox_inches='tight', pad_inches=0.05,\n",
        "                        dpi=200)\n",
        "        if self.show_fig:\n",
        "            plt.show()\n",
        "        else:\n",
        "            plt.close('all')\n",
        "            \n",
        "    \n",
        "    def _KmeansAcc(self, x, y, c):\n",
        "        kmeans = KMeans(n_clusters=c, n_jobs=-1, n_init=10, random_state=12)\n",
        "        \n",
        "        # Increase the value of rindex from c*2, its is only a demo!\n",
        "        rindex = random.sample(range(0, y.shape[0]), max(min(self.nodes, y.shape[0]), c*2))\n",
        "        outs = self.AE(self.mel(x[rindex]), training=False)\n",
        "        y_pred = kmeans.fit_predict(outs)\n",
        "       \n",
        "        acc = accuracy(y[rindex], y_pred)\n",
        "        nmi = normalized_mutual_info_score(y[rindex], y_pred)\n",
        "        ari = adjusted_rand_score(y[rindex], y_pred)\n",
        "        return (acc, nmi, ari, outs, y[rindex])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6AXWbOTV3k9P"
      },
      "source": [
        "class AEGenerator(tf.keras.utils.Sequence):\n",
        "    \"\"\" Recieves X and y \n",
        "        Performes pairwise matching with a batch size \n",
        "        both pairs are generated from the input X and y\n",
        " \n",
        "        X         : Input, shape=(bs, farme_size, 1)\n",
        "        y         : Output shape=(bs, ) or (bs, 1)\n",
        "        dist      : Distance parameter for AE\n",
        "        noise     : Noise for augmentation, shape=(bs, farme_size, 1)\n",
        "        scale     : The maximum limit of noise that would be mixed with X\n",
        "        show_logs : Show random selection errors\n",
        "        gt        : Ground truth of the actual class.\n",
        "                    Use only when show_logs is True\n",
        "    \"\"\"\n",
        " \n",
        "    def __init__(self, iX, iy, dist, noise, batch_size=64, \n",
        "                 scale=0.4, show_logs=False, gt=False,\n",
        "                 # MEL layer setup\n",
        "                 mel_scale='spec', setup=0, output_channel=1,\n",
        "                 ):\n",
        "        self.batch_size  = batch_size\n",
        "        self.noise       = noise\n",
        "\n",
        "        self.mel = MEL(scale=mel_scale, setup=setup, \n",
        "                       output_channel=output_channel)\n",
        "        \n",
        "        # inputX : mel, inputA : raw-audio, inputY : labels \n",
        "        #self.inputX = self.mel(iX).numpy()\n",
        "        self.inputA = iX\n",
        "        self.inputY = iy\n",
        "        self.output_shape = self.mel(self.inputA[:1, ...]).numpy().shape\n",
        "\n",
        "        if self.inputY.shape[-1] == 1:\n",
        "            self.inputY = np.squeeze(self.inputY)\n",
        "        if self.inputA.shape[-1] == 1:\n",
        "            self.inputA = np.squeeze(self.inputA)\n",
        "        if self.noise.shape[-1] == 1:\n",
        "            self.noise = np.squeeze(self.noise)\n",
        "        \n",
        "        self.total       = len(iX)\n",
        "        self.dist        = dist\n",
        "        self.scale       = scale\n",
        "        self.gt          = gt\n",
        " \n",
        "        # Generates label : indexes_where_label_found\n",
        "        ulabels = np.unique(self.inputY)\n",
        "        self.class_index = dict([(label, list(np.where(self.inputY == label)[0])) \\\n",
        "                                 for label in ulabels])  \n",
        " \n",
        "        self.indexes     = np.arange(self.total)\n",
        "        self.total_batch = self.total // self.batch_size \n",
        "        self.classes = len(np.unique(self.inputY))\n",
        " \n",
        "        self.rand = random.Random(12)\n",
        "        random.seed(12)\n",
        " \n",
        "        # Show logs\n",
        "        self.show_logs = show_logs\n",
        "        self.log = {}\n",
        "        self.on_epoch_end()\n",
        "    \n",
        "    def _augment(self, wav_indices):\n",
        "        \"\"\"Needs to be modified\n",
        "           Input_dim : [SR, 1]\n",
        "        \"\"\"\n",
        "        noise_id = np.random.randint(low=0, high=self.noise.shape[0], \n",
        "                                     size=self.batch_size)\n",
        "        scale = np.random.uniform(low=0, high=self.scale,\n",
        "                                  size=(self.batch_size, 1))\n",
        "        \n",
        "        real = np.multiply(self.inputA[wav_indices], (1-scale)) \n",
        "        jitter = np.multiply(self.noise[noise_id], scale)\n",
        "        wav = real + jitter\n",
        "\n",
        "        return wav\n",
        "    \n",
        "    def __len__(self):\n",
        "        \"\"\" Denotes the number of batches per epoch \"\"\"\n",
        "        return self.total_batch\n",
        " \n",
        "    def _print_logs(self):\n",
        "        for a, b in self.log.items():\n",
        "            print(a, ':', b)\n",
        "    \n",
        "    def on_epoch_end(self):\n",
        "        \"\"\" Updates indexes after each epoch \"\"\"\n",
        "        np.random.shuffle(self.indexes)\n",
        "        if self.show_logs:\n",
        "            self._print_logs()    \n",
        "            self.log['canNotLink_error'] = 0\n",
        "            self.log['time'] = 0\n",
        " \n",
        "    def _make_choice(self, p=None):\n",
        "        w = [0.5, 0.5]\n",
        "        if p is not None:\n",
        "            w = [p, 1-p]\n",
        "        return self.rand.choices([True, False], weights=w)[0]\n",
        " \n",
        "    def __getitem__(self, batch_index):\n",
        "        \"\"\" Generate one batch of data \"\"\"\n",
        "        idx, y = self._genIndexes(batch_index)\n",
        "        wavs = self.inputA[idx]\n",
        "        \n",
        "        if self.scale > 0:\n",
        "            aug_idx = self.rand.sample(range(2*self.batch_size), \n",
        "                                       k=self.batch_size)\n",
        "            wavs[aug_idx] = self._augment(idx[aug_idx])\n",
        "\n",
        "        rets = self.mel(np.expand_dims(wavs, axis=-1)).numpy()\n",
        "        return [rets[:self.batch_size], rets[self.batch_size:]], y\n",
        " \n",
        "    def _genIndexes(self, index):\n",
        "        idx = np.zeros((2, self.batch_size), dtype=np.int32)\n",
        "        idx[0] = self.indexes[index*self.batch_size:(index+1)*self.batch_size]\n",
        "        constraint = np.zeros(self.batch_size, dtype=np.int32)\n",
        " \n",
        "        for i in range(self.batch_size):\n",
        "            plabel = self.inputY[idx[0][i]]\n",
        "            take = idx[0][i]\n",
        " \n",
        "            # Generating must-links\n",
        "            if self._make_choice():\n",
        "                take = self.rand.choice(self.class_index[plabel])\n",
        "            # Generating can-not links\n",
        "            else:\n",
        "                # Taking random pair\n",
        "                while self.inputY[take] == plabel:\n",
        "                    take = self.rand.choice(self.indexes)\n",
        "                constraint[i] = self.dist\n",
        " \n",
        "                if self.show_logs and self.gt[idx[0][i]] == self.gt[take]:\n",
        "                    self.log['canNotLink_error'] += 1\n",
        "\n",
        "            idx[1][i] = take\n",
        "\n",
        "        return idx.reshape((2*self.batch_size)), constraint"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MHFcm10B3ljj"
      },
      "source": [
        "# Implementation of the distance layer\n",
        "# We are assuming that the dimenstion of both of the embedders are \n",
        "# already subtracted\n",
        "class Distance(Layer):\n",
        "    def __init__(self, ):\n",
        "        super(Distance, self).__init__()\n",
        "        \n",
        "    def call(self, inputs):\n",
        "        return tf.expand_dims(\n",
        "               tf.math.sqrt(\n",
        "               tf.math.reduce_sum(\n",
        "               tf.math.square(inputs), axis=1)), axis=-1)\n",
        "\n",
        "\n",
        "# Implementation of the AutoEmbedder\n",
        "def buildAE(input_shape, dims, alpha=1, dis=100, topmodel='MobileNet',\n",
        "            scale=None, setup=0, dropout=0.001, decay=False):\n",
        "    \n",
        "    inp1 = Input(input_shape, name='input1')\n",
        "    inp2 = Input(input_shape, name='input2')\n",
        "\n",
        "\n",
        "    topmodel = eval(f\"{topmodel}(input_shape={input_shape},\" + \n",
        "                        \"include_top=False, weights='imagenet')\")\n",
        "    \n",
        "    m = Sequential([topmodel,\n",
        "                    Flatten(), \n",
        "                    Dense(dims)], name='AutoEmbedder')\n",
        "    #m = SincNet(input_shape, dims, softmax=True)\n",
        "\n",
        "    out1 = m(inp1)\n",
        "    out2 = m(inp2)\n",
        "\n",
        "    # Subtracting output pairs\n",
        "    out = Subtract()([out1, out2])\n",
        "    # Calculating distance\n",
        "    out = Distance()(out)\n",
        "    # The thresholded ReLU layer\n",
        "    out = ReLU(max_value=dis)(out)\n",
        "\n",
        "    # Initializing model\n",
        "    model = Model([inp1, inp2], out, name=\"AE_train\")\n",
        "    # Using default adam optimizer with mean square error\n",
        "    if decay:\n",
        "        print('Decay loaded')\n",
        "        lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n",
        "                       initial_learning_rate=0.1, decay_steps=50,\n",
        "                       decay_rate=0.0005)\n",
        "\n",
        "    # Optimizer default:\n",
        "    # tf.keras.optimizers.Adam(learning_rate=0.0007 if decay is False else lr_schedule)\n",
        "    model.compile(optimizer= tf.keras.optimizers.Adam(learning_rate=0.0005), \n",
        "                   loss='mse')\n",
        "    AE = Model(inp1, out1, name=\"AutoEmbedder\")    \n",
        "\n",
        "    return model, AE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EJB48DZS3nTN"
      },
      "source": [
        "base_dir = './logs/DenseNet_ASR/'\n",
        "model, AE = None, None\n",
        "\n",
        "def load_AE(path, input_shape):\n",
        "    model = tf.keras.models.load_model(path)\n",
        "    print('Model loaded')\n",
        "    inp = Input(input_shape)\n",
        "    out = model.get_layer('AutoEmbedder')(inp)\n",
        "    AE = Model(inp, out)\n",
        "    print('AE loaded')\n",
        "    return model, AE\n",
        " \n",
        " \n",
        "def trainModel(model, X, cy, y, noise, db, spkrs, clim, \n",
        "               alpha=0.25, dist=100, dims=3, setup=2,\n",
        "               batch_size=64, scale='spec', epochs=3000, decay=False,\n",
        "               load_pretrain=False, frame=0.2, seconds=16,\n",
        "               add_mel=True, aug_scale=0, base_dir=base_dir, \n",
        "               pretrain_path=None):\n",
        "    '''\n",
        "    This would check if previously ran model exists, and load it in model and \n",
        "    AE. This function would return MyLogger callback object and the starting \n",
        "    epoch for the model. Necessary savepath parameters would be declared in this\n",
        "    function as parameter.\n",
        " \n",
        "    Possible Naming Scheme: Model_dist_dims_setup_scale_frame_[T,L,B]_MIN_SPKRs_CLIM\n",
        "    T : TIMIT\n",
        "    L : LibriSpeech\n",
        "    B : Bengali ASR\n",
        "    dist  : Distance between clusters (known as AutoEmbedder alpha parameter)\n",
        "    dims  : The output dimention of AutoEmbedder\n",
        "    frame : Frame size (input frame size)\n",
        "    scale : The scaling feature of speech (spectrogram, mfcc)\n",
        "    SPKRS : Number of speakers\n",
        "    CLIM  : Maximum Inter-cluster-linkage\n",
        "    MIN   : Minutes of speech per-speaker\n",
        "    '''\n",
        "    global AE\n",
        "    tf.keras.backend.clear_session()\n",
        "    nmodel = f\"{model}{alpha}\" if model == 'MobileNet' else f\"{model}\"\n",
        "    pname = f\"{nmodel}_A{dist}_D{dims}_{setup}_{scale}_F{frame}_{db}_T{seconds}_S{spkrs}_L{clim}\"\n",
        "    \n",
        "    loaded = False\n",
        "    save_dir = os.path.join(base_dir, pname)\n",
        "    \n",
        "    dg = AEGenerator(iX=X, iy=cy, dist=dist, noise=noise, \n",
        "                     batch_size=batch_size, scale=aug_scale, \n",
        "                     show_logs=False, gt=y, \n",
        "                     output_channel=1 if model == 'MobileNet' else 3)\n",
        "\n",
        "    #if not os.path.exists(save_dir):\n",
        "    #    os.makedirs(save_dir)\n",
        "    if os.path.exists(os.path.join(save_dir, 'model')):\n",
        "        try:\n",
        "            print('Previous model found, loading model')\n",
        "            model, AE = load_AE(os.path.join(save_dir, 'model'), \n",
        "                                dg.output_shape[1:])\n",
        "            loaded = True\n",
        "        except Exception as e:\n",
        "            print(str(e))\n",
        "            pass\n",
        "    \n",
        "    if not loaded and load_pretrain and pretrain_path is not None:\n",
        "        print('Loading pre-trained model')\n",
        "        model, AE = load_AE(pretrain_path, dg.output_shape[1:])\n",
        "        loaded = True\n",
        " \n",
        "\n",
        "    if not loaded:\n",
        "        print('Initializing new model')\n",
        "        model, AE = buildAE(dg.output_shape[1:], \n",
        "                            dims=dims, setup=setup, \n",
        "                            alpha=alpha, scale=scale, dis=100, \n",
        "                            decay=decay, topmodel=model)\n",
        "        loaded = True\n",
        "    \n",
        "\n",
        "    print(X.shape, cy.shape, y.shape)\n",
        "    log = MyLogger([(X, cy), (X, y)], n=5, plot=True, AE=AE,\n",
        "                   savepath=save_dir, \n",
        "                   save_model=True,\n",
        "                   scatter=False,\n",
        "                   val_train=True,\n",
        "                   show_fig=False,\n",
        "                   nodes=len(np.unique(y))*2,\n",
        "                   output_channel= (1 if nmodel.startswith('MobileNet') else 3))\n",
        "\n",
        "    model.summary()\n",
        "    model.get_layer('AutoEmbedder').summary()\n",
        "    #time.sleep(1)\n",
        "    #AE.summary()\n",
        "    \n",
        "    print('Starting model training')\n",
        "    model.fit(dg, epochs=epochs+1, verbose=0, callbacks=[log], \n",
        "              initial_epoch=log.start_epoch, use_multiprocessing=False,\n",
        "              workers=4, max_queue_size=10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qez1AKkb3rPo"
      },
      "source": [
        "frame = 0.2\n",
        "seconds = 10\n",
        "spkrs = 25\n",
        "clim = 5\n",
        "dbi = 1\n",
        "db = ['TIMIT', 'LIBRI', 'ASR']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cM_4WTYp7kwy"
      },
      "source": [
        "X, y, fi, noise = loadData(db[dbi], frame=frame, seconds=seconds)\n",
        " \n",
        "print(X.shape, y.shape)\n",
        " \n",
        "if spkrs is not None and spkrs != 0:\n",
        "    X, y = speakerFilter(X, y, spkrs)\n",
        " \n",
        "spkrs = len(np.unique(y))\n",
        "cy = pairwiseRelations(y, max_lim=clim)\n",
        " \n",
        "print(spkrs, clim, seconds, len(np.unique(cy)))\n",
        "print(X.shape, y.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hmaa4sjE7mKV"
      },
      "source": [
        "base_dir = f'./logs/'\n",
        "\n",
        "trainModel(model='DenseNet121', alpha=1, dist=100, dims=12, setup=0, \n",
        "           scale='spec', db=db[dbi][0], spkrs=spkrs, clim=clim, \n",
        "           X=X, cy=cy, y=y, noise=noise, batch_size=128,\n",
        "           load_pretrain=False, frame=frame, epochs=1200,\n",
        "           seconds=seconds, decay=False, add_mel=False,\n",
        "           aug_scale=0.07, base_dir=base_dir, \n",
        "           pretrain_path=None)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}