Wave-U-net.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Wave-U-net.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyNZnuv8jiKS20LkFXDggXAL",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/satvik-venkatesh/Wave-U-net-TF2/blob/main/Wave-U-net.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6b9K3J4ZEdqV"
      },
      "source": [
        "import tensorflow as tf\r\n",
        "import numpy as np\r\n",
        "from tensorflow.keras.layers import Layer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qAyowzcelZdt"
      },
      "source": [
        "# Custom layers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LFkCcACAQJA3"
      },
      "source": [
        "class AudioClipLayer(Layer):\r\n",
        "\r\n",
        "    def __init__(self, **kwargs):\r\n",
        "        '''Initializes the instance attributes'''\r\n",
        "        super(AudioClipLayer, self).__init__(**kwargs)\r\n",
        "\r\n",
        "    def build(self, input_shape):\r\n",
        "        '''Create the state of the layer (weights)'''\r\n",
        "        # initialize the weights\r\n",
        "        pass\r\n",
        "        \r\n",
        "    def call(self, inputs, training):\r\n",
        "        '''Defines the computation from inputs to outputs'''\r\n",
        "        if training:\r\n",
        "            return inputs\r\n",
        "        else:\r\n",
        "            return tf.maximum(tf.minimum(inputs, 1.0), -1.0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Woi9w83c93fG"
      },
      "source": [
        "# Learned Interpolation layer\r\n",
        "\r\n",
        "class InterpolationLayer(Layer):\r\n",
        "\r\n",
        "    def __init__(self, padding = \"valid\", **kwargs):\r\n",
        "        '''Initializes the instance attributes'''\r\n",
        "        super(InterpolationLayer, self).__init__(**kwargs)\r\n",
        "        self.padding = padding\r\n",
        "\r\n",
        "    def build(self, input_shape):\r\n",
        "        '''Create the state of the layer (weights)'''\r\n",
        "        self.features = input_shape.as_list()[3]\r\n",
        "\r\n",
        "        # initialize the weights\r\n",
        "        w_init = tf.random_normal_initializer()\r\n",
        "        self.w = tf.Variable(name=\"kernel\",\r\n",
        "            initial_value=w_init(shape=(self.features, ),\r\n",
        "                                 dtype='float32'),\r\n",
        "            trainable=True)\r\n",
        "\r\n",
        "    def call(self, inputs):\r\n",
        "        '''Defines the computation from inputs to outputs'''\r\n",
        "\r\n",
        "        w_scaled = tf.math.sigmoid(self.w)\r\n",
        "\r\n",
        "        counter_w = 1 - w_scaled\r\n",
        "\r\n",
        "        conv_weights = tf.expand_dims(tf.concat([tf.expand_dims(tf.linalg.diag(w_scaled), axis=0), tf.expand_dims(tf.linalg.diag(counter_w), axis=0)], axis=0), axis=0)\r\n",
        "\r\n",
        "        intermediate_vals = tf.nn.conv2d(inputs, conv_weights, strides=[1,1,1,1], padding=self.padding.upper())\r\n",
        "\r\n",
        "        intermediate_vals = tf.transpose(intermediate_vals, [2, 0, 1, 3])\r\n",
        "        out = tf.transpose(inputs, [2, 0, 1, 3])\r\n",
        "        \r\n",
        "        num_entries = out.shape.as_list()[0]\r\n",
        "        out = tf.concat([out, intermediate_vals], axis=0)\r\n",
        "\r\n",
        "        indices = list()\r\n",
        "\r\n",
        "        # num_outputs = 2*num_entries - 1\r\n",
        "        num_outputs = (2*num_entries - 1) if self.padding == \"valid\" else 2*num_entries\r\n",
        "\r\n",
        "        for idx in range(num_outputs):\r\n",
        "            if idx % 2 == 0:\r\n",
        "                indices.append(idx // 2)\r\n",
        "            else:\r\n",
        "                indices.append(num_entries + idx//2)\r\n",
        "        out = tf.gather(out, indices)\r\n",
        "        current_layer = tf.transpose(out, [1, 2, 0, 3])\r\n",
        "\r\n",
        "        return current_layer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9onhDb0lTCQp"
      },
      "source": [
        "class CropLayer(Layer):\r\n",
        "    def __init__(self, x2, match_feature_dim=True, **kwargs):\r\n",
        "        '''Initializes the instance attributes'''\r\n",
        "        super(CropLayer, self).__init__(**kwargs)\r\n",
        "        self.match_feature_dim = match_feature_dim\r\n",
        "        self.x2 = x2\r\n",
        "\r\n",
        "    def build(self, input_shape):\r\n",
        "        '''Create the state of the layer (weights)'''\r\n",
        "        # initialize the weights\r\n",
        "        pass\r\n",
        "        \r\n",
        "    def call(self, inputs):\r\n",
        "        '''Defines the computation from inputs to outputs'''\r\n",
        "        if self.x2 is None:\r\n",
        "            return inputs\r\n",
        "\r\n",
        "        inputs = self.crop(inputs, self.x2.shape.as_list(), self.match_feature_dim)\r\n",
        "        return inputs\r\n",
        "\r\n",
        "    def crop(self, tensor, target_shape, match_feature_dim=True):\r\n",
        "        '''\r\n",
        "        Crops a 3D tensor [batch_size, width, channels] along the width axes to a target shape.\r\n",
        "        Performs a centre crop. If the dimension difference is uneven, crop last dimensions first.\r\n",
        "        :param tensor: 4D tensor [batch_size, width, height, channels] that should be cropped. \r\n",
        "        :param target_shape: Target shape (4D tensor) that the tensor should be cropped to\r\n",
        "        :return: Cropped tensor\r\n",
        "        '''\r\n",
        "        shape = np.array(tensor.shape.as_list())\r\n",
        "\r\n",
        "        ddif = shape[1] - target_shape[1]\r\n",
        "\r\n",
        "        if (ddif % 2 != 0):\r\n",
        "            print(\"WARNING: Cropping with uneven number of extra entries on one side\")\r\n",
        "        # assert diff[1] >= 0 # Only positive difference allowed\r\n",
        "        if ddif == 0:\r\n",
        "            return tensor\r\n",
        "        crop_start = ddif // 2\r\n",
        "        crop_end = ddif - crop_start\r\n",
        "\r\n",
        "        return tensor[:,crop_start:-crop_end,:]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1qlZHtgp-6aF"
      },
      "source": [
        "class IndependentOutputLayer(Layer):\r\n",
        "\r\n",
        "    def __init__(self, source_names, num_channels, filter_width, padding=\"valid\", **kwargs):\r\n",
        "        '''Initializes the instance attributes'''\r\n",
        "        super(IndependentOutputLayer, self).__init__(**kwargs)\r\n",
        "        self.source_names = source_names\r\n",
        "        self.num_channels = num_channels\r\n",
        "        self.filter_width = filter_width\r\n",
        "        self.padding = padding\r\n",
        "\r\n",
        "        self.conv1a = tf.keras.layers.Conv1D(self.num_channels, self.filter_width, padding= self.padding)\r\n",
        "\r\n",
        "\r\n",
        "    def build(self, input_shape):\r\n",
        "        '''Create the state of the layer (weights)'''\r\n",
        "        pass\r\n",
        "        \r\n",
        "    def call(self, inputs, training):\r\n",
        "        '''Defines the computation from inputs to outputs'''\r\n",
        "        outputs = {}\r\n",
        "        for name in self.source_names:\r\n",
        "            out = self.conv1a(inputs)\r\n",
        "            outputs[name] = out\r\n",
        "        \r\n",
        "        return outputs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oxLPPJ1Lasvn"
      },
      "source": [
        "class DiffOutputLayer(Layer):\r\n",
        "\r\n",
        "    def __init__(self, source_names, num_channels, filter_width, padding=\"valid\", **kwargs):\r\n",
        "        '''Initializes the instance attributes'''\r\n",
        "        super(DiffOutputLayer, self).__init__(**kwargs)\r\n",
        "        self.source_names = source_names\r\n",
        "        self.num_channels = num_channels\r\n",
        "        self.filter_width = filter_width\r\n",
        "        self.padding = padding\r\n",
        "\r\n",
        "        self.conv1a = tf.keras.layers.Conv1D(self.num_channels, self.filter_width, padding= self.padding)\r\n",
        "\r\n",
        "\r\n",
        "    def build(self, input_shape):\r\n",
        "        '''Create the state of the layer (weights)'''\r\n",
        "        pass\r\n",
        "        \r\n",
        "    def call(self, inputs, training):\r\n",
        "        '''Defines the computation from inputs to outputs'''\r\n",
        "        outputs = {}\r\n",
        "        sum_source = 0\r\n",
        "        for name in self.source_names[:-1]:\r\n",
        "            out = self.conv1a(inputs[0])\r\n",
        "            out = AudioClipLayer()(out)\r\n",
        "            outputs[name] = out\r\n",
        "            sum_source = sum_source + out\r\n",
        "        \r\n",
        "        last_source = CropLayer(sum_source)(inputs[1]) - sum_source\r\n",
        "        last_source = AudioClipLayer()(last_source)\r\n",
        "\r\n",
        "        outputs[self.source_names[-1]] = last_source\r\n",
        "\r\n",
        "        return outputs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UV6T2ttZlg7Q"
      },
      "source": [
        "# Define the Network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_4MYB1pafSfI"
      },
      "source": [
        "def wave_u_net(num_initial_filters = 24, num_layers = 12, kernel_size = 15, merge_filter_size = 5, \r\n",
        "               source_names = [\"bass\", \"drums\", \"other\", \"vocals\"], num_channels = 1, output_filter_size = 1,\r\n",
        "               padding = \"same\", input_size = 16384 * 4, context = False, upsampling_type = \"learned\",\r\n",
        "               output_activation = \"linear\", output_type = \"difference\"):\r\n",
        "  \r\n",
        "  # `enc_outputs` stores the downsampled outputs to re-use during upsampling.\r\n",
        "  enc_outputs = []\r\n",
        "\r\n",
        "  # `raw_input` is the input to the network\r\n",
        "  raw_input = tf.keras.layers.Input(shape=(input_size, num_channels),name=\"raw_input\")\r\n",
        "  X = raw_input\r\n",
        "  inp = raw_input\r\n",
        "\r\n",
        "  # Down sampling\r\n",
        "  for i in range(num_layers):\r\n",
        "    X = tf.keras.layers.Conv1D(filters=num_initial_filters + (num_initial_filters * i),\r\n",
        "                          kernel_size=kernel_size,strides=1,\r\n",
        "                          padding=padding, name=\"Down_Conv_\"+str(i))(X)\r\n",
        "    X = tf.keras.layers.LeakyReLU(name=\"Down_Conv_Activ_\"+str(i))(X)\r\n",
        "\r\n",
        "    enc_outputs.append(X)\r\n",
        "\r\n",
        "    X = tf.keras.layers.Lambda(lambda x: x[:,::2,:], name=\"Decimate_\"+str(i))(X)\r\n",
        "\r\n",
        "\r\n",
        "  X = tf.keras.layers.Conv1D(filters=num_initial_filters + (num_initial_filters * num_layers),\r\n",
        "                          kernel_size=kernel_size,strides=1,\r\n",
        "                          padding=padding, name=\"Down_Conv_\"+str(num_layers))(X)\r\n",
        "  X = tf.keras.layers.LeakyReLU(name=\"Down_Conv_Activ_\"+str(num_layers))(X)\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "  # Up sampling\r\n",
        "  for i in range(num_layers):\r\n",
        "    X = tf.keras.layers.Lambda(lambda x: tf.expand_dims(x, axis=1), name=\"exp_dims_\"+str(i))(X)\r\n",
        "    \r\n",
        "    if upsampling_type == \"learned\":\r\n",
        "      X = InterpolationLayer(name=\"IntPol_\"+str(i), padding=padding)(X)\r\n",
        "\r\n",
        "    else:\r\n",
        "      if context:\r\n",
        "        X = tf.keras.layers.Lambda(lambda x: tf.image.resize(x, [1, x.shape.as_list()[2] * 2 - 1]), name=\"bilinear_interpol_\"+str(i))(X)\r\n",
        "        # current_layer = tf.image.resize_bilinear(current_layer, [1, current_layer.get_shape().as_list()[2] * 2 - 1], align_corners=True)\r\n",
        "      else:\r\n",
        "        X = tf.keras.layers.Lambda(lambda x: tf.image.resize(x, [1, x.shape.as_list()[2] * 2]), name=\"bilinear_interpol_\"+str(i))(X)\r\n",
        "        # current_layer = tf.image.resize_bilinear(current_layer, [1, current_layer.get_shape().as_list()[2]*2]) # out = in + in - 1\r\n",
        "\r\n",
        "\r\n",
        "    X = tf.keras.layers.Lambda(lambda x: tf.squeeze(x, axis=1), name=\"sq_dims_\"+str(i))(X)\r\n",
        "    \r\n",
        "    c_layer = CropLayer(X, False, name=\"crop_layer_\"+str(i))(enc_outputs[-i-1])\r\n",
        "    X = tf.keras.layers.Concatenate(axis=2, name=\"concatenate_\"+str(i))([X, c_layer]) \r\n",
        "\r\n",
        "\r\n",
        "    X = tf.keras.layers.Conv1D(filters=num_initial_filters + (num_initial_filters * (num_layers - i - 1)),\r\n",
        "                            kernel_size=merge_filter_size,strides=1,\r\n",
        "                            padding=padding, name=\"Up_Conv_\"+str(i))(X)\r\n",
        "    X = tf.keras.layers.LeakyReLU(name=\"Up_Conv_Activ_\"+str(i))(X)\r\n",
        "\r\n",
        "\r\n",
        "  c_layer = CropLayer(X, False, name=\"crop_layer_\"+str(num_layers))(inp)\r\n",
        "  X = tf.keras.layers.Concatenate(axis=2, name=\"concatenate_\"+str(num_layers))([X, c_layer]) \r\n",
        "  X = AudioClipLayer(name=\"audio_clip_\"+str(0))(X)\r\n",
        "\r\n",
        "  if output_type == \"direct\":\r\n",
        "    X = IndependentOutputLayer(source_names, num_channels, output_filter_size, padding=padding, name=\"independent_out\")(X)\r\n",
        "\r\n",
        "  else:\r\n",
        "    # Difference Output\r\n",
        "    cropped_input = CropLayer(X, False, name=\"crop_layer_\"+str(num_layers+1))(inp)\r\n",
        "    X = DiffOutputLayer(source_names, num_channels, output_filter_size, padding=padding, name=\"diff_out\")([X, cropped_input])\r\n",
        "\r\n",
        "  o = X\r\n",
        "  model = tf.keras.Model(inputs=raw_input, outputs=o)\r\n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cRqjX5C5APAX"
      },
      "source": [
        "# Parameters for the Wave-U-net\r\n",
        "\r\n",
        "params = {\r\n",
        "  \"num_initial_filters\": 24,\r\n",
        "  \"num_layers\": 12,\r\n",
        "  \"kernel_size\": 15,\r\n",
        "  \"merge_filter_size\": 5,\r\n",
        "  \"source_names\": [\"bass\", \"drums\", \"other\", \"vocals\"],\r\n",
        "  \"num_channels\": 2,\r\n",
        "  \"output_filter_size\": 1,\r\n",
        "  \"padding\": \"valid\",\r\n",
        "  \"input_size\": 147443,\r\n",
        "  \"context\": True,\r\n",
        "  \"upsampling_type\": \"learned\",         # \"learned\" or \"linear\"\r\n",
        "  \"output_activation\": \"linear\",        # \"linear\" or \"tanh\"\r\n",
        "  \"output_type\": \"difference\",          # \"direct\" or \"difference\" \r\n",
        "}"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tPc2_kxuf_xE"
      },
      "source": [
        "m = wave_u_net(**params)"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DEZbrek4mqwg",
        "outputId": "319cb542-34b5-4ee3-e2f2-fdd05823a557"
      },
      "source": [
        "m.summary()"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"model_5\"\n",
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            "raw_input (InputLayer)          [(None, 147443, 2)]  0                                            \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_0 (Conv1D)            (None, 147429, 24)   744         raw_input[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_0 (LeakyReLU)   (None, 147429, 24)   0           Down_Conv_0[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_0 (Lambda)             (None, 73715, 24)    0           Down_Conv_Activ_0[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_1 (Conv1D)            (None, 73701, 48)    17328       Decimate_0[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_1 (LeakyReLU)   (None, 73701, 48)    0           Down_Conv_1[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_1 (Lambda)             (None, 36851, 48)    0           Down_Conv_Activ_1[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_2 (Conv1D)            (None, 36837, 72)    51912       Decimate_1[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_2 (LeakyReLU)   (None, 36837, 72)    0           Down_Conv_2[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_2 (Lambda)             (None, 18419, 72)    0           Down_Conv_Activ_2[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_3 (Conv1D)            (None, 18405, 96)    103776      Decimate_2[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_3 (LeakyReLU)   (None, 18405, 96)    0           Down_Conv_3[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_3 (Lambda)             (None, 9203, 96)     0           Down_Conv_Activ_3[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_4 (Conv1D)            (None, 9189, 120)    172920      Decimate_3[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_4 (LeakyReLU)   (None, 9189, 120)    0           Down_Conv_4[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_4 (Lambda)             (None, 4595, 120)    0           Down_Conv_Activ_4[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_5 (Conv1D)            (None, 4581, 144)    259344      Decimate_4[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_5 (LeakyReLU)   (None, 4581, 144)    0           Down_Conv_5[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_5 (Lambda)             (None, 2291, 144)    0           Down_Conv_Activ_5[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_6 (Conv1D)            (None, 2277, 168)    363048      Decimate_5[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_6 (LeakyReLU)   (None, 2277, 168)    0           Down_Conv_6[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_6 (Lambda)             (None, 1139, 168)    0           Down_Conv_Activ_6[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_7 (Conv1D)            (None, 1125, 192)    484032      Decimate_6[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_7 (LeakyReLU)   (None, 1125, 192)    0           Down_Conv_7[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_7 (Lambda)             (None, 563, 192)     0           Down_Conv_Activ_7[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_8 (Conv1D)            (None, 549, 216)     622296      Decimate_7[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_8 (LeakyReLU)   (None, 549, 216)     0           Down_Conv_8[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_8 (Lambda)             (None, 275, 216)     0           Down_Conv_Activ_8[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_9 (Conv1D)            (None, 261, 240)     777840      Decimate_8[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_9 (LeakyReLU)   (None, 261, 240)     0           Down_Conv_9[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_9 (Lambda)             (None, 131, 240)     0           Down_Conv_Activ_9[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_10 (Conv1D)           (None, 117, 264)     950664      Decimate_9[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_10 (LeakyReLU)  (None, 117, 264)     0           Down_Conv_10[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_10 (Lambda)            (None, 59, 264)      0           Down_Conv_Activ_10[0][0]         \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_11 (Conv1D)           (None, 45, 288)      1140768     Decimate_10[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_11 (LeakyReLU)  (None, 45, 288)      0           Down_Conv_11[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Decimate_11 (Lambda)            (None, 23, 288)      0           Down_Conv_Activ_11[0][0]         \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_12 (Conv1D)           (None, 9, 312)       1348152     Decimate_11[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "Down_Conv_Activ_12 (LeakyReLU)  (None, 9, 312)       0           Down_Conv_12[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_0 (Lambda)             (None, 1, 9, 312)    0           Down_Conv_Activ_12[0][0]         \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_0 (InterpolationLayer)   (None, 1, 17, 312)   312         exp_dims_0[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_0 (Lambda)              (None, 17, 312)      0           IntPol_0[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_0 (CropLayer)        (None, 17, 288)      0           Down_Conv_Activ_11[0][0]         \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_0 (Concatenate)     (None, 17, 600)      0           sq_dims_0[0][0]                  \n",
            "                                                                 crop_layer_0[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_0 (Conv1D)              (None, 13, 288)      864288      concatenate_0[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_0 (LeakyReLU)     (None, 13, 288)      0           Up_Conv_0[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_1 (Lambda)             (None, 1, 13, 288)   0           Up_Conv_Activ_0[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_1 (InterpolationLayer)   (None, 1, 25, 288)   288         exp_dims_1[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_1 (Lambda)              (None, 25, 288)      0           IntPol_1[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_1 (CropLayer)        (None, 25, 264)      0           Down_Conv_Activ_10[0][0]         \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_1 (Concatenate)     (None, 25, 552)      0           sq_dims_1[0][0]                  \n",
            "                                                                 crop_layer_1[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_1 (Conv1D)              (None, 21, 264)      728904      concatenate_1[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_1 (LeakyReLU)     (None, 21, 264)      0           Up_Conv_1[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_2 (Lambda)             (None, 1, 21, 264)   0           Up_Conv_Activ_1[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_2 (InterpolationLayer)   (None, 1, 41, 264)   264         exp_dims_2[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_2 (Lambda)              (None, 41, 264)      0           IntPol_2[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_2 (CropLayer)        (None, 41, 240)      0           Down_Conv_Activ_9[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_2 (Concatenate)     (None, 41, 504)      0           sq_dims_2[0][0]                  \n",
            "                                                                 crop_layer_2[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_2 (Conv1D)              (None, 37, 240)      605040      concatenate_2[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_2 (LeakyReLU)     (None, 37, 240)      0           Up_Conv_2[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_3 (Lambda)             (None, 1, 37, 240)   0           Up_Conv_Activ_2[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_3 (InterpolationLayer)   (None, 1, 73, 240)   240         exp_dims_3[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_3 (Lambda)              (None, 73, 240)      0           IntPol_3[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_3 (CropLayer)        (None, 73, 216)      0           Down_Conv_Activ_8[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_3 (Concatenate)     (None, 73, 456)      0           sq_dims_3[0][0]                  \n",
            "                                                                 crop_layer_3[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_3 (Conv1D)              (None, 69, 216)      492696      concatenate_3[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_3 (LeakyReLU)     (None, 69, 216)      0           Up_Conv_3[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_4 (Lambda)             (None, 1, 69, 216)   0           Up_Conv_Activ_3[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_4 (InterpolationLayer)   (None, 1, 137, 216)  216         exp_dims_4[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_4 (Lambda)              (None, 137, 216)     0           IntPol_4[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_4 (CropLayer)        (None, 137, 192)     0           Down_Conv_Activ_7[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_4 (Concatenate)     (None, 137, 408)     0           sq_dims_4[0][0]                  \n",
            "                                                                 crop_layer_4[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_4 (Conv1D)              (None, 133, 192)     391872      concatenate_4[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_4 (LeakyReLU)     (None, 133, 192)     0           Up_Conv_4[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_5 (Lambda)             (None, 1, 133, 192)  0           Up_Conv_Activ_4[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_5 (InterpolationLayer)   (None, 1, 265, 192)  192         exp_dims_5[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_5 (Lambda)              (None, 265, 192)     0           IntPol_5[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_5 (CropLayer)        (None, 265, 168)     0           Down_Conv_Activ_6[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_5 (Concatenate)     (None, 265, 360)     0           sq_dims_5[0][0]                  \n",
            "                                                                 crop_layer_5[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_5 (Conv1D)              (None, 261, 168)     302568      concatenate_5[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_5 (LeakyReLU)     (None, 261, 168)     0           Up_Conv_5[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_6 (Lambda)             (None, 1, 261, 168)  0           Up_Conv_Activ_5[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_6 (InterpolationLayer)   (None, 1, 521, 168)  168         exp_dims_6[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_6 (Lambda)              (None, 521, 168)     0           IntPol_6[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_6 (CropLayer)        (None, 521, 144)     0           Down_Conv_Activ_5[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_6 (Concatenate)     (None, 521, 312)     0           sq_dims_6[0][0]                  \n",
            "                                                                 crop_layer_6[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_6 (Conv1D)              (None, 517, 144)     224784      concatenate_6[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_6 (LeakyReLU)     (None, 517, 144)     0           Up_Conv_6[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_7 (Lambda)             (None, 1, 517, 144)  0           Up_Conv_Activ_6[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_7 (InterpolationLayer)   (None, 1, 1033, 144) 144         exp_dims_7[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_7 (Lambda)              (None, 1033, 144)    0           IntPol_7[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_7 (CropLayer)        (None, 1033, 120)    0           Down_Conv_Activ_4[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_7 (Concatenate)     (None, 1033, 264)    0           sq_dims_7[0][0]                  \n",
            "                                                                 crop_layer_7[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_7 (Conv1D)              (None, 1029, 120)    158520      concatenate_7[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_7 (LeakyReLU)     (None, 1029, 120)    0           Up_Conv_7[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_8 (Lambda)             (None, 1, 1029, 120) 0           Up_Conv_Activ_7[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_8 (InterpolationLayer)   (None, 1, 2057, 120) 120         exp_dims_8[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_8 (Lambda)              (None, 2057, 120)    0           IntPol_8[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_8 (CropLayer)        (None, 2057, 96)     0           Down_Conv_Activ_3[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_8 (Concatenate)     (None, 2057, 216)    0           sq_dims_8[0][0]                  \n",
            "                                                                 crop_layer_8[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_8 (Conv1D)              (None, 2053, 96)     103776      concatenate_8[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_8 (LeakyReLU)     (None, 2053, 96)     0           Up_Conv_8[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_9 (Lambda)             (None, 1, 2053, 96)  0           Up_Conv_Activ_8[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_9 (InterpolationLayer)   (None, 1, 4105, 96)  96          exp_dims_9[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_9 (Lambda)              (None, 4105, 96)     0           IntPol_9[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_9 (CropLayer)        (None, 4105, 72)     0           Down_Conv_Activ_2[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_9 (Concatenate)     (None, 4105, 168)    0           sq_dims_9[0][0]                  \n",
            "                                                                 crop_layer_9[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_9 (Conv1D)              (None, 4101, 72)     60552       concatenate_9[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_9 (LeakyReLU)     (None, 4101, 72)     0           Up_Conv_9[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_10 (Lambda)            (None, 1, 4101, 72)  0           Up_Conv_Activ_9[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_10 (InterpolationLayer)  (None, 1, 8201, 72)  72          exp_dims_10[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_10 (Lambda)             (None, 8201, 72)     0           IntPol_10[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_10 (CropLayer)       (None, 8201, 48)     0           Down_Conv_Activ_1[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_10 (Concatenate)    (None, 8201, 120)    0           sq_dims_10[0][0]                 \n",
            "                                                                 crop_layer_10[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_10 (Conv1D)             (None, 8197, 48)     28848       concatenate_10[0][0]             \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_10 (LeakyReLU)    (None, 8197, 48)     0           Up_Conv_10[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "exp_dims_11 (Lambda)            (None, 1, 8197, 48)  0           Up_Conv_Activ_10[0][0]           \n",
            "__________________________________________________________________________________________________\n",
            "IntPol_11 (InterpolationLayer)  (None, 1, 16393, 48) 48          exp_dims_11[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "sq_dims_11 (Lambda)             (None, 16393, 48)    0           IntPol_11[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_11 (CropLayer)       (None, 16393, 24)    0           Down_Conv_Activ_0[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_11 (Concatenate)    (None, 16393, 72)    0           sq_dims_11[0][0]                 \n",
            "                                                                 crop_layer_11[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_11 (Conv1D)             (None, 16389, 24)    8664        concatenate_11[0][0]             \n",
            "__________________________________________________________________________________________________\n",
            "Up_Conv_Activ_11 (LeakyReLU)    (None, 16389, 24)    0           Up_Conv_11[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_12 (CropLayer)       (None, 16389, 2)     0           raw_input[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "concatenate_12 (Concatenate)    (None, 16389, 26)    0           Up_Conv_Activ_11[0][0]           \n",
            "                                                                 crop_layer_12[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "audio_clip_0 (AudioClipLayer)   (None, 16389, 26)    0           concatenate_12[0][0]             \n",
            "__________________________________________________________________________________________________\n",
            "crop_layer_13 (CropLayer)       (None, 16389, 2)     0           raw_input[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "diff_out (DiffOutputLayer)      {'bass': (None, 1638 54          audio_clip_0[0][0]               \n",
            "                                                                 crop_layer_13[0][0]              \n",
            "==================================================================================================\n",
            "Total params: 10,265,550\n",
            "Trainable params: 10,265,550\n",
            "Non-trainable params: 0\n",
            "__________________________________________________________________________________________________\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IUQ3cplMm8Hf",
        "outputId": "cbe17e09-1e49-468a-9d46-87522ad855ee"
      },
      "source": [
        "m.output"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'bass': <KerasTensor: shape=(None, 16389, 2) dtype=float32 (created by layer 'diff_out')>,\n",
              " 'drums': <KerasTensor: shape=(None, 16389, 2) dtype=float32 (created by layer 'diff_out')>,\n",
              " 'other': <KerasTensor: shape=(None, 16389, 2) dtype=float32 (created by layer 'diff_out')>,\n",
              " 'vocals': <KerasTensor: shape=(None, 16389, 2) dtype=float32 (created by layer 'diff_out')>}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WfjeTMngEI-u"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_2xXpQWolnyY"
      },
      "source": [
        "# Other utility functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HxhDdi8SfXRS"
      },
      "source": [
        "def get_padding(shape, num_layers=12, filter_size=15, input_filter_size=15, output_filter_size=1, merge_filter_size=5, num_channels=1, context = True):\r\n",
        "    '''\r\n",
        "    Note that this function is not used within the Wave-U-net. \r\n",
        "    But it is useful to calculate the required amounts of padding along \r\n",
        "    each axis of the input and output, so that the Unet works and has the \r\n",
        "    given shape as output shape.\r\n",
        "\r\n",
        "    :param shape: Desired output shape \r\n",
        "    :return: Input_shape, output_shape, where each is a list [batch_size, time_steps, channels]\r\n",
        "    '''\r\n",
        "\r\n",
        "    if context:\r\n",
        "        # Check if desired shape is possible as output shape - go from output shape towards lowest-res feature map\r\n",
        "        rem = float(shape[1]) # Cut off batch size number and channel\r\n",
        "\r\n",
        "        # Output filter size\r\n",
        "        rem = rem - output_filter_size + 1\r\n",
        "\r\n",
        "        # Upsampling blocks\r\n",
        "        for i in range(num_layers):\r\n",
        "            rem = rem + merge_filter_size - 1\r\n",
        "            rem = (rem + 1.) / 2.# out = in + in - 1 <=> in = (out+1)/\r\n",
        "\r\n",
        "        # Round resulting feature map dimensions up to nearest integer\r\n",
        "        x = np.asarray(np.ceil(rem),dtype=np.int64)\r\n",
        "        assert(x >= 2)\r\n",
        "\r\n",
        "        # Compute input and output shapes based on lowest-res feature map\r\n",
        "        output_shape = x\r\n",
        "        input_shape = x\r\n",
        "\r\n",
        "        # Extra conv\r\n",
        "        input_shape = input_shape + filter_size - 1\r\n",
        "\r\n",
        "        # Go from centre feature map through up- and downsampling blocks\r\n",
        "        for i in range(num_layers):\r\n",
        "            output_shape = 2*output_shape - 1 #Upsampling\r\n",
        "            output_shape = output_shape - merge_filter_size + 1 # Conv\r\n",
        "\r\n",
        "            input_shape = 2*input_shape - 1 # Decimation\r\n",
        "            if i < num_layers - 1:\r\n",
        "                input_shape = input_shape + filter_size - 1 # Conv\r\n",
        "            else:\r\n",
        "                input_shape = input_shape + input_filter_size - 1\r\n",
        "\r\n",
        "        # Output filters\r\n",
        "        output_shape = output_shape - output_filter_size + 1\r\n",
        "\r\n",
        "        input_shape = np.concatenate([[shape[0]], [input_shape], [num_channels]])\r\n",
        "        output_shape = np.concatenate([[shape[0]], [output_shape], [num_channels]])\r\n",
        "\r\n",
        "        return input_shape, output_shape\r\n",
        "    else:\r\n",
        "        return [shape[0], shape[1], num_channels], [shape[0], shape[1], num_channels]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ptkj9jCVid0T",
        "outputId": "de3cce66-82be-4718-e26e-fd74c1f13b5a"
      },
      "source": [
        "get_padding((16, 16389, 1))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(array([    16, 147443,      1]), array([   16, 16389,     1]))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    }
  ]
}