Deep_Complex_U_net.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Deep Complex U-net.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyN1JEWsZ0ccqcw9CDzVShqY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mehdihosseinimoghadam/Signal-Processing/blob/main/Deep_Complex_U_net.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "NZPxUl5yrPV1"
      },
      "outputs": [],
      "source": [
        "from torch import nn\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class CConv2d(nn.Module):\n",
        "  def __init__(self, in_channels, out_channels, **kwargs):\n",
        "    super(CConv2d, self).__init__()\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "\n",
        "\n",
        "    self.re_conv = nn.Conv2d(self.in_channels, self.out_channels, **kwargs)\n",
        "    self.im_conv = nn.Conv2d(self.in_channels, self.out_channels, **kwargs)\n",
        "\n",
        "    nn.init.xavier_uniform_(self.re_conv.weight)\n",
        "    nn.init.xavier_uniform_(self.im_conv.weight)\n",
        "\n",
        "  def forward(self, x):  \n",
        "    x_re = x[..., 0]\n",
        "    x_im = x[..., 1]\n",
        "\n",
        "    out_re = self.re_conv[x_re] - self.im_conv(x_im)\n",
        "    out_im = self.re_conv[x_im] + self.im_conv(x_re)\n",
        "\n",
        "    out = torch.cat([out_re, out_im], -1) \n",
        "\n",
        "    return out\n"
      ],
      "metadata": {
        "id": "-ACZSe4arcbG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class CConvTrans2d(nn.Module):\n",
        "  def __init__(self, in_channels, out_channels, **kwargs):\n",
        "    super(CConvTrans2d, self).__init__()\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "\n",
        "\n",
        "    self.re_Tconv = nn.ConvTranspose2d(self.in_channels, self.out_channels, **kwargs)\n",
        "    self.im_Tconv = nn.ConvTranspose2d(self.in_channels, self.out_channels, **kwargs)\n",
        "\n",
        "    nn.init.xavier_uniform_(self.re_Tconv.weight)\n",
        "    nn.init.xavier_uniform_(self.im_Tconv.weight)\n",
        "\n",
        "\n",
        "  def forward(self, x):  \n",
        "    x_re = x[..., 0]\n",
        "    x_im = x[..., 1]\n",
        "\n",
        "    out_re = self.re_Tconv[x_re] - self.im_Tconv(x_im)\n",
        "    out_im = self.re_Tconv[x_im] + self.im_Tconv(x_re)\n",
        "\n",
        "    out = torch.cat([out_re, out_im], -1) \n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "rAFkSP1BvOFn"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class CBatchnorm(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, **kwargs):\n",
        "        super(CBatchnorm, self).__init__()\n",
        "        self.in_channels = in_channels\n",
        "\n",
        "        self.re_batch = nn.BatchNorm2d(in_channels)\n",
        "        self.im_batch = nn.BatchNorm2d(in_channels)\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        x_re = x[..., 0]\n",
        "        x_im = x[..., 1]\n",
        "\n",
        "        out_re =  self.re_batch(x_re)\n",
        "        out_im =  self.re_batch(x_im)\n",
        "\n",
        "\n",
        "        out = torch.cat([out_re, out_im], -1) \n",
        "\n",
        "        return out\n",
        "\n"
      ],
      "metadata": {
        "id": "TolGeNh3xO0z"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class CconvBlock(nn.Module):\n",
        "  def __init__(self, in_channels, out_channels, kernel, stride, padding):\n",
        "    super(CconvBlock, self).__init()\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "    self.kernel = kernel\n",
        "    self.stride = stride\n",
        "    self.padding = padding\n",
        "\n",
        "    self.CConv2d = CConv2d(self.in_channels, self.out_channels, self.kernel, self.stride , self.padding)\n",
        "    self.CBatchnorm = CBatchnorm(self.out_channels)\n",
        "    self.leaky_relu = nn.LeakyReLU()\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    conved = self.CConv2d(x)\n",
        "    normed = self.CBatchnorm(conved)\n",
        "    activated =  self.leaky_relu(normed)\n",
        "\n",
        "    return activated\n",
        "\n"
      ],
      "metadata": {
        "id": "rhkZ62u9zI7Z"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class CConvTransBlock(nn.Module):\n",
        "  def __init__(self, in_channels, out_channels, kernel, stride, padding, last_layer=False):\n",
        "    super(CConvTransBlock, self).__init()\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "    self.kernel = kernel\n",
        "    self.stride = stride\n",
        "    self.padding = padding\n",
        "    self.last_layer = last_layer\n",
        "\n",
        "    self.CConvTrans2d = CConvTrans2d(self.in_channels, self.out_channels, self.kernel, self.stride , self.padding)\n",
        "    self.CBatchnorm = CBatchnorm(self.out_channels)\n",
        "    self.leaky_relu = nn.LeakyReLU()\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    conved =  self.CConvTrans2d(x)\n",
        "\n",
        "    if not self.last_layer:\n",
        "        normed = self.CBatchnorm(conved)\n",
        "        activated =  self.leaky_relu(normed)\n",
        "        return activated\n",
        "    else:\n",
        "        m_phase = conved/(torch.abs(conved)+1e-8)  \n",
        "        m_mag = torch.tanh(torch.abs(conved))\n",
        "        out = m_phase * m_mag\n",
        "        return out  \n"
      ],
      "metadata": {
        "id": "ne9KCvIm6nOv"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Encoder(nn.Module):\n",
        "  def __init__(self):\n",
        "      super(Encoder, self).__init__()\n",
        "\n",
        "      self.CconvBlock0 = CconvBlock(filter_size=(7,5), stride_size=(2,2), in_channels=1, out_channels=45, padding=(0,0))\n",
        "      self.CconvBlock1 = CconvBlock(filter_size=(7,5), stride_size=(2,2), in_channels=45, out_channels=90, padding=(0,0))\n",
        "      self.CconvBlock2 = CconvBlock(filter_size=(5,3), stride_size=(2,2), in_channels=90, out_channels=90, padding=(0,0))\n",
        "      self.CconvBlock3 = CconvBlock(filter_size=(5,3), stride_size=(2,2), in_channels=90, out_channels=90, padding=(0,0))\n",
        "      self.CconvBlock4 = CconvBlock(filter_size=(5,3), stride_size=(2,1), in_channels=90, out_channels=90, padding=(0,0))\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "      ccb0 = self.CconvBlock0(x)\n",
        "      ccb1 = self.CconvBlock1(ccb0) \n",
        "      ccb2 = self.CconvBlock2(ccb1)        \n",
        "      ccb3 = self.CconvBlock3(ccb2)        \n",
        "      ccb4 = self.CconvBlock4(ccb3)\n",
        "\n",
        "      return [ccb0, ccb1, ccb2, ccb3, ccb4]"
      ],
      "metadata": {
        "id": "C97FsXBhEZJz"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Decoder(nn.Module):\n",
        "  def __init__(self, n_fft=64, hop_length=16):\n",
        "    super(Decoder, self).__init__()\n",
        "    self.n_fft = n_fft\n",
        "    self.hop_length = hop_length\n",
        "\n",
        "\n",
        "    self.CConvTransBlock0 = CConvTransBlock(filter_size=(5,3), stride_size=(2,1), in_channels=90, out_channels=90, output_padding=(0,0), padding=(0,0))\n",
        "    self.CConvTransBlock1 = CConvTransBlock(filter_size=(5,3), stride_size=(2,2), in_channels=180, out_channels=90, output_padding=(0,0), padding=(0,0))\n",
        "    self.CConvTransBlock2 = CConvTransBlock(filter_size=(5,3), stride_size=(2,2), in_channels=180, out_channels=90, output_padding=(0,0), padding=(0,0))\n",
        "    self.CConvTransBlock3 = CConvTransBlock(filter_size=(7,5), stride_size=(2,2), in_channels=180, out_channels=45, output_padding=(0,0), padding=(0,0))\n",
        "    self.CConvTransBlock4 = CConvTransBlock(filter_size=(7,5), stride_size=(2,2), in_channels=90, output_padding=(0,1), padding=(0,0),\n",
        "                              out_channels=1, last_layer=True)\n",
        "    \n",
        "\n",
        "  def forward(self, x0, x, is_istft=True):\n",
        "\n",
        "    \n",
        "        cctb0 = self.CConvTransBlock0(x[-1])\n",
        "        # skip-connection\n",
        "        c0 = torch.cat((cctb0, x[-2]), dim=1)\n",
        "        \n",
        "        cctb1 = self.CConvTransBlock1(c0)\n",
        "        c1 = torch.cat((cctb1, x[-3]), dim=1)\n",
        "        \n",
        "        cctb2 = self.CConvTransBlock2(c1)\n",
        "        c2 = torch.cat((cctb2, x[-4]), dim=1)\n",
        "        \n",
        "        cctb3 = self.CConvTransBlock3(c2)\n",
        "        c3 = torch.cat((cctb3, x[-5]), dim=1)\n",
        "        \n",
        "        cctb4 = self.CConvTransBlock4(c3)\n",
        "\n",
        "\n",
        "        output = cctb4 * x0\n",
        "        if is_istft:\n",
        "            output = torch.squeeze(output, 1)\n",
        "            output = torch.istft(output, n_fft=self.n_fft, hop_length=self.hop_length, normalized=True)\n",
        "        \n",
        "        return output\n"
      ],
      "metadata": {
        "id": "mrYxCS-5Id71"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Model(nn.Module):\n",
        "  def __init__(self, encoder_out, n_fft=64, hop_length=16):\n",
        "    super(Model, self).__init__()\n",
        "    self.encoder_out = encoder_out\n",
        "    self.n_fft = n_fft\n",
        "    self.hop_length = hop_length\n",
        "\n",
        "    self.Encoder = Encoder()\n",
        "    self.Decoder = Decoder(self.encoder_out, self.n_fft, self.hop_length)\n",
        "\n",
        "  def forward(self, x):\n",
        "      encoded = self.Encoder(x)\n",
        "      decoded = self.Decoder(x, encoded) \n",
        "      return decoded"
      ],
      "metadata": {
        "id": "72uMa77DNI3y"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}