FastCycle-GAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IMPORTING LIBRARIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FFT LAYER DEFINITION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft_forward(image,w,b,hparameters1):\n",
    "    \n",
    "    '''\n",
    "    FUNCTION DESCRIPTION:\n",
    "    This function is used to transform and convolute in spectral domain using the FFT approach. The purpose of this\n",
    "    function is to excecute a fast convolution that is less time consuming than the spatial convolution given that in the\n",
    "    spectral domain a convolution is a simple function product, in this case between the filter and the image.\n",
    "    \n",
    "    PARAMETER DESCRIPTION:\n",
    "    - image: a matrix representation of an image to process\n",
    "    - w: the weight to apply to the convolution\n",
    "    - b: the weight to add to the convoluted image\n",
    "    - hparameters1: **********************\n",
    "    '''\n",
    "    \n",
    "    # Defining image shape\n",
    "    length = image.shape[1]\n",
    "    padding = (length - 3)//2\n",
    "    filter = np.pad(w.reshape(3,3), (padding,padding))\n",
    "    image = image.reshape(length,length)\n",
    "    \n",
    "    # For even values of length\n",
    "    if length % 2 == 0:\n",
    "        length = length-1\n",
    "        image = image[:length, :length]\n",
    "        limit = length//2+1\n",
    "        \n",
    "        # Transforming image and filter using FFT\n",
    "        fft_image = np.fft.fft2(image)\n",
    "        fft_filter = np.fft.fft2(filter)\n",
    "        \n",
    "        # Convolution in Spectral Domain\n",
    "        conv = fft_image * fft_filter\n",
    "        \n",
    "        # Returning to Spatial Domain\n",
    "        conv_image = np.fft.ifft2(conv)\n",
    "        \n",
    "        # Transposing image\n",
    "        y = np.zeros((length,length))\n",
    "        y[1:limit, 1:limit] = conv_img.real[limit:, limit:]\n",
    "        y[1:limit, limit:] = conv_img.real[limit:, 1:limit]\n",
    "        y[limit:, 1:limit] = conv_img.real[1:limit, limit:]\n",
    "        y[limit:, limit:] = conv_img.real[1:limit, 1:limit]\n",
    "\n",
    "        # Returning function\n",
    "        y = y[1:length, 1:length] + b\n",
    "    \n",
    "    # For odd values of length\n",
    "    else:\n",
    "        limit = length//2+1\n",
    "        \n",
    "        # Transforming image and filter using FFT\n",
    "        fft_image = np.fft.fft2(image)\n",
    "        fft_filter = np.fft.fft2(filter)\n",
    "        \n",
    "        # Convolution in Spectral Domain\n",
    "        conv = fft_image * fft_filter\n",
    "        \n",
    "        # Returning to Spatial Domain\n",
    "        conv_image = np.fft.ifft2(conv)\n",
    "        \n",
    "        # Transposing image\n",
    "        y = np.zeros((length,length))\n",
    "        y[1:limit , 1:limit] = conv_image.real[limit:, limit:]\n",
    "        y[1:limit , limit:] = conv_image.real[limit:, 1:limit]\n",
    "        y[limit: , 1:limit] = conv_image.real[1:limit, limit:]\n",
    "        y[limit: , limit:] = conv_image.real[1:limit, 1:limit]\n",
    "        \n",
    "        # Returning function\n",
    "        y = y[1:length-1, 1:length-1] + b\n",
    "        \n",
    "    return zeros"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SPECTRAL POOLING LAYER DEFINITION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'torch'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-20-355ab13610ea>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;31m# import torch\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnn\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mnn\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mautograd\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mFunction\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mmath\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodules\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mutils\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0m_pair\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'torch'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Function\n",
    "import math\n",
    "from torch.nn.modules.utils import _pair\n",
    "\n",
    "def _spectral_crop(input, oheight, owidth):\n",
    "\n",
    "    cutoff_freq_h = math.ceil(oheight / 2)\n",
    "    cutoff_freq_w = math.ceil(owidth / 2)\n",
    "\n",
    "    if oheight % 2 == 1:\n",
    "        if owidth % 2 == 1:\n",
    "            top_left = input[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            top_right = input[:, :, :cutoff_freq_h, -(cutoff_freq_w-1):]\n",
    "            bottom_left = input[:, :, -(cutoff_freq_h-1):, :cutoff_freq_w]\n",
    "            bottom_right = input[:, :, -(cutoff_freq_h-1):, -(cutoff_freq_w-1):]\n",
    "        else:\n",
    "            top_left = input[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            top_right = input[:, :, :cutoff_freq_h, -cutoff_freq_w:]\n",
    "            bottom_left = input[:, :, -(cutoff_freq_h-1):, :cutoff_freq_w]\n",
    "            bottom_right = input[:, :, -(cutoff_freq_h-1):, -cutoff_freq_w:]\n",
    "    else:\n",
    "        if owidth % 2 == 1:\n",
    "            top_left = input[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            top_right = input[:, :, :cutoff_freq_h, -(cutoff_freq_w-1):]\n",
    "            bottom_left = input[:, :, -cutoff_freq_h:, :cutoff_freq_w]\n",
    "            bottom_right = input[:, :, -cutoff_freq_h:, -(cutoff_freq_w-1):]\n",
    "        else:\n",
    "            top_left = input[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            top_right = input[:, :, :cutoff_freq_h, -cutoff_freq_w:]\n",
    "            bottom_left = input[:, :, -cutoff_freq_h:, :cutoff_freq_w]\n",
    "            bottom_right = input[:, :, -cutoff_freq_h:, -cutoff_freq_w:]\n",
    "\n",
    "    top_combined = torch.cat((top_left, top_right), dim=-1)\n",
    "    bottom_combined = torch.cat((bottom_left, bottom_right), dim=-1)\n",
    "    all_together = torch.cat((top_combined, bottom_combined), dim=-2)\n",
    "\n",
    "    return all_together\n",
    "\n",
    "def _spectral_pad(input, output, oheight, owidth):\n",
    "    cutoff_freq_h = math.ceil(oheight / 2)\n",
    "    cutoff_freq_w = math.ceil(owidth / 2)\n",
    "\n",
    "    pad = torch.zeros_like(input)\n",
    "\n",
    "    if oheight % 2 == 1:\n",
    "        if owidth % 2 == 1:\n",
    "            pad[:, :, :cutoff_freq_h, :cutoff_freq_w] = output[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            pad[:, :, :cutoff_freq_h, -(cutoff_freq_w-1):] = output[:, :, :cutoff_freq_h, -(cutoff_freq_w-1):]\n",
    "            pad[:, :, -(cutoff_freq_h-1):, :cutoff_freq_w] = output[:, :, -(cutoff_freq_h-1):, :cutoff_freq_w]\n",
    "            pad[:, :, -(cutoff_freq_h-1):, -(cutoff_freq_w-1):] = output[:, :, -(cutoff_freq_h-1):, -(cutoff_freq_w-1):]\n",
    "        else:\n",
    "            pad[:, :, :cutoff_freq_h, :cutoff_freq_w] = output[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            pad[:, :, :cutoff_freq_h, -cutoff_freq_w:] = output[:, :, :cutoff_freq_h, -cutoff_freq_w:]\n",
    "            pad[:, :, -(cutoff_freq_h-1):, :cutoff_freq_w] = output[:, :, -(cutoff_freq_h-1):, :cutoff_freq_w]\n",
    "            pad[:, :, -(cutoff_freq_h-1):, -cutoff_freq_w:] = output[:, :, -(cutoff_freq_h-1):, -cutoff_freq_w:]\n",
    "    else:\n",
    "        if owidth % 2 == 1:\n",
    "            pad[:, :, :cutoff_freq_h, :cutoff_freq_w] = output[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            pad[:, :, :cutoff_freq_h, -(cutoff_freq_w-1):] = output[:, :, :cutoff_freq_h, -(cutoff_freq_w-1):]\n",
    "            pad[:, :, -cutoff_freq_h:, :cutoff_freq_w] = output[:, :, -cutoff_freq_h:, :cutoff_freq_w]\n",
    "            pad[:, :, -cutoff_freq_h:, -(cutoff_freq_w-1):] = output[:, :, -cutoff_freq_h:, -(cutoff_freq_w-1):]\n",
    "        else:\n",
    "            pad[:, :, :cutoff_freq_h, :cutoff_freq_w] = output[:, :, :cutoff_freq_h, :cutoff_freq_w]\n",
    "            pad[:, :, :cutoff_freq_h, -cutoff_freq_w:] = output[:, :, :cutoff_freq_h, -cutoff_freq_w:]\n",
    "            pad[:, :, -cutoff_freq_h:, :cutoff_freq_w] = output[:, :, -cutoff_freq_h:, :cutoff_freq_w]\n",
    "            pad[:, :, -cutoff_freq_h:, -cutoff_freq_w:] = output[:, :, -cutoff_freq_h:, -cutoff_freq_w:]\t\n",
    "\n",
    "    return pad\n",
    "\n",
    "def DiscreteHartleyTransform(input):\n",
    "    fft = np.fft.fft2(input, 2, normalized=True, onesided=False)\n",
    "    dht = fft[:, :, :, :, -2] - fft[:, :, :, :, -1]\n",
    "    return dht\n",
    "\n",
    "class SpectralPoolingFunction(Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, input, oheight, owidth):\n",
    "        ctx.oh = oheight\n",
    "        ctx.ow = owidth\n",
    "        ctx.save_for_backward(input)\n",
    "\n",
    "        # Hartley transform by RFFT\n",
    "        dht = DiscreteHartleyTransform(input)\n",
    "\n",
    "        # frequency cropping\n",
    "        all_together = _spectral_crop(dht, oheight, owidth)\n",
    "        \n",
    "        # inverse Hartley transform\n",
    "        dht = DiscreteHartleyTransform(all_together)\n",
    "        return dht\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        input, = ctx.saved_variables\n",
    "\n",
    "        # Hartley transform by RFFT\n",
    "        dht = DiscreteHartleyTransform(grad_output)\n",
    "        # frequency padding\n",
    "        grad_input = _spectral_pad(input, dht, ctx.oh, ctx.ow)\n",
    "        # inverse Hartley transform\n",
    "        grad_input = DiscreteHartleyTransform(grad_input)\n",
    "        return grad_input, None, None\n",
    "\n",
    "class SpectralPool2d(nn.Module):\n",
    "    def __init__(self, scale_factor):\n",
    "        super(SpectralPool2d, self).__init__()\n",
    "        self.scale_factor = _pair(scale_factor)\n",
    "    def forward(self, input):\n",
    "        H, W = input.size(-2), input.size(-1)\n",
    "        h, w = math.ceil(H*self.scale_factor[0]), math.ceil(W*self.scale_factor[1])\n",
    "        return SpectralPoolingFunction.apply(input, h, w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft_model(image, weights, hparameters1, hparameters2):\n",
    "    fft = relu(fft_forward(img, weights[0], weights[1], hparameters1))\n",
    "    pool = pool_forward(fft.reshape(1,223,223,1), hparameters2)[0]\n",
    "    fft1 = relu(fft_forward(pool, weights[2], weights[3], hparameters1))\n",
    "    pool1 = pool_forward(fft1.reshape(1,fft1.shape[1],fft1.shape[1],1), hparameters2)[0]\n",
    "    fft2 = relu(fft_forward(pool1, weights[4], weights[5], hparameters1))\n",
    "    pool2 = pool_forward(fft2.reshape(1,fft2.shape[1],fft2.shape[1],1), hparameters2)[0]\n",
    "    flatten = pool2.reshape(1,pool2.shape[1]*pool2.shape[2]*pool2.shape[3])\n",
    "    neural_net = relu(np.dot(flatten, weights[6]) + weights[7])\n",
    "    ans = sigmoid(np.dot(neural_net, weights[8]) + weights[9])\n",
    "    \n",
    "    tup = (ans, fft, pool, fft1, pool1, fft2, pool2)\n",
    "    return tup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hparameters1 = {\n",
    "    'stride':1,\n",
    "    'pad':0\n",
    "}\n",
    "hparameters2 = {\n",
    "    'stride':2,\n",
    "    'f':2\n",
    "}\n",
    "weights = model.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fft_out = []\n",
    "\n",
    "start = time.time()\n",
    "for i in range(len(x_test)):\n",
    "    img = cv2.resize(x_test[i],(225,225)).reshape(1,225,225,1)\n",
    "    output1 = fft_model(img, weights, hparameters1, hparameters2)\n",
    "    if output1[0] >= 0.5:\n",
    "        fft_out.append(1)\n",
    "    else:\n",
    "        fft_out.append(0)\n",
    "\n",
    "end = time.time()\n",
    "print('Time taken by FFT to predict class for 253 images is % seconds', %end - start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Activation(x,function_name='leaky_relu'):\n",
    "    \n",
    "    def lrelu(x, leak=0.2, name=function_name, alt_relu_impl=False):\n",
    "\n",
    "        with tf.variable_scope(name):\n",
    "            if alt_relu_impl:\n",
    "                f1 = 0.5 * (1 + leak)\n",
    "                f2 = 0.5 * (1 - leak)\n",
    "                # lrelu = 1/2 * (1 + leak) * x + 1/2 * (1 - leak) * |x|\n",
    "                return f1 * x + f2 * abs(x)\n",
    "            else:\n",
    "                return tf.maximum(x, leak * x)\n",
    "\n",
    "    def relu(x, function_name):\n",
    "        return np.maximum(0, x)\n",
    "\n",
    "    def sigmoid(x, function_name):\n",
    "        return 1 / (1 + np.exp(-x))\n",
    "\n",
    "\n",
    "    \n",
    "def InstanceNormalization(x):\n",
    "\n",
    "    with tf.variable_scope(\"instance_norm\"):\n",
    "        epsilon = 1e-5\n",
    "        \n",
    "        mean, var = tf.nn.moments(x,\n",
    "                                  [1, 2],\n",
    "                                  keep_dims=True)\n",
    "        \n",
    "        scale = tf.get_variable('scale',\n",
    "                                [x.get_shape()[-1]], \n",
    "                                initializer=tf.truncated_normal_initializer(mean=1.0, stddev=0.02))\n",
    "        \n",
    "        offset = tf.get_variable('offset',\n",
    "                                 [x.get_shape()[-1]],\n",
    "                                 initializer=tf.constant_initializer(0.0))\n",
    "        \n",
    "        out = scale*tf.div(x-mean, tf.sqrt(var+epsilon)) + offset\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### COLOR TRANSFORMATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lab2rgb(self, L, AB):\n",
    "        \"\"\"Convert an Lab tensor image to a RGB numpy output\n",
    "        Parameters:\n",
    "            L  (1-channel tensor array): L channel images (range: [-1, 1], torch tensor array)\n",
    "            AB (2-channel tensor array):  ab channel images (range: [-1, 1], torch tensor array)\n",
    "        Returns:\n",
    "            rgb (RGB numpy image): rgb output images  (range: [0, 255], numpy array)\n",
    "        \"\"\"\n",
    "        AB2 = AB * 110.0\n",
    "        L2 = (L + 1.0) * 50.0\n",
    "        Lab = torch.cat([L2, AB2], dim=1)\n",
    "        Lab = Lab[0].data.cpu().float().numpy()\n",
    "        Lab = np.transpose(Lab.astype(np.float64), (1, 2, 0))\n",
    "        rgb = color.lab2rgb(Lab) * 255\n",
    "        return rgb\n",
    "\n",
    "    \n",
    "def rgb2lab (inputColor):\n",
    "\n",
    "    num = 0\n",
    "    RGB = [0, 0, 0]\n",
    "\n",
    "    for value in inputColor :\n",
    "        value = float(value) / 255\n",
    "\n",
    "        if value > 0.04045 :\n",
    "            value = ( ( value + 0.055 ) / 1.055 ) ** 2.4\n",
    "        else :\n",
    "            value = value / 12.92\n",
    "\n",
    "        RGB[num] = value * 100\n",
    "        num = num + 1\n",
    "\n",
    "    XYZ = [0, 0, 0,]\n",
    "\n",
    "    X = RGB [0] * 0.4124 + RGB [1] * 0.3576 + RGB [2] * 0.1805\n",
    "    Y = RGB [0] * 0.2126 + RGB [1] * 0.7152 + RGB [2] * 0.0722\n",
    "    Z = RGB [0] * 0.0193 + RGB [1] * 0.1192 + RGB [2] * 0.9505\n",
    "    XYZ[ 0 ] = round( X, 4 )\n",
    "    XYZ[ 1 ] = round( Y, 4 )\n",
    "    XYZ[ 2 ] = round( Z, 4 )\n",
    "\n",
    "    XYZ[ 0 ] = float( XYZ[ 0 ] ) / 95.047         # ref_X =  95.047   Observer= 2°, Illuminant= D65\n",
    "    XYZ[ 1 ] = float( XYZ[ 1 ] ) / 100.0          # ref_Y = 100.000\n",
    "    XYZ[ 2 ] = float( XYZ[ 2 ] ) / 108.883        # ref_Z = 108.883\n",
    "\n",
    "    num = 0\n",
    "    for value in XYZ :\n",
    "\n",
    "        if value > 0.008856 :\n",
    "            value = value ** ( 0.3333333333333333 )\n",
    "        else :\n",
    "            value = ( 7.787 * value ) + ( 16 / 116 )\n",
    "\n",
    "        XYZ[num] = value\n",
    "        num = num + 1\n",
    "\n",
    "    Lab = [0, 0, 0]\n",
    "\n",
    "    L = ( 116 * XYZ[ 1 ] ) - 16\n",
    "    a = 500 * ( XYZ[ 0 ] - XYZ[ 1 ] )\n",
    "    b = 200 * ( XYZ[ 1 ] - XYZ[ 2 ] )\n",
    "\n",
    "    Lab [ 0 ] = round( L, 4 )\n",
    "    Lab [ 1 ] = round( a, 4 )\n",
    "    Lab [ 2 ] = round( b, 4 )\n",
    "\n",
    "    return Lab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'filename' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-5-06c9308c5cd4>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mskimage\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mio\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcolor\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[0mrgb\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mio\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mimread\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfilename\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[0mlab\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcolor\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrgb2lab\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mrgb\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[0mrgb\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcolor\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlab2rgb\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mlab\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'filename' is not defined"
     ]
    }
   ],
   "source": [
    "from skimage import io, color\n",
    "rgb = io.imread(img)\n",
    "lab = color.rgb2lab(rgb)\n",
    "rgb = color.lab2rgb(lab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MODEL DEFINITION\n",
    "\n",
    "We define the model hyperparameters and image preprocessing. Given that we are going to work with an image colorization problem, the most suitable color space for this task is LAB instead of RGB.<br>LAB is a color space defined by three dimensions: Particularly L for luminosity, describes the values between black and white; A is used to specify the values between green and red; And B indicates the values between yellow and blue.<br>The values of L go from 0 to 100, 0 being the description for black and 100 for white. While the A and B values go from negative values (green and blue respectively) to positive values (red and yellow).<br>This model particularly will have as its input the L dimension of the image and its main goal is to generate the other two dimensions of the color space.<br>The images that with which the model will be fed will consis of a 256x256 matrix with only one channel (the L dimension of the LAB space).<p>For the model definition we proceed to define the hyperparameters and the corresponding layers of the model. In this case we will be working firstly with a ResNet-like architecture for both Generators (A and B) being in an autoencoder arrangement, the convolution and deconvolution layers are performed using a FFT. On the other hand, the Discriminators will have a four-layer FFT convolution architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The learning rate is set to 0.000000 given the number of epochs.\n"
     ]
    }
   ],
   "source": [
    "# Image preprocessing\n",
    "img_height = 256\n",
    "img_width = 256\n",
    "img_size = img_height * img_width\n",
    "image = cv2.resize(x_test[15],(img_height,img_width)).reshape(1,img_height,img_width,1) \n",
    "for im in image:\n",
    "    rgb = io.imread(im)\n",
    "    lab = color.rgb2lab(rgb)\n",
    "    print(image)\n",
    "\n",
    "# Defining model hyperparameters\n",
    "batch_size = 50\n",
    "pool_size = 100\n",
    "epoch=1000\n",
    "gen_features = 32\n",
    "disc_features = 64\n",
    "learning_rate = lr_schedule(epoch)\n",
    "\n",
    "# Setting Learning Rate for different number of Epochs\n",
    "def lr_schedule(epoch):\n",
    "    lr = 1e-3\n",
    "    if epoch > 180:\n",
    "        lr *= 0.5e-3\n",
    "    elif epoch > 160:\n",
    "        lr *= 1e-3\n",
    "    elif epoch > 120:\n",
    "        lr *= 1e-2\n",
    "    elif epoch > 80:\n",
    "        lr *= 1e-1\n",
    "    print('The learning rate is set to %f given the number of epochs.' %(lr))\n",
    "    return lr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RESNET BLOCK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic ResNet Building Block\n",
    "def FastConvolutionLayer(inputs,\n",
    "                 num_filters = 16,\n",
    "                 kernel_size = 3,\n",
    "                 strides = 1,\n",
    "                 activation ='relu',\n",
    "                 batch_normalization = False,\n",
    "                 conv_first=False):\n",
    "    \n",
    "    conv = fft_forward(num_filters,w,b,hparameters)\n",
    "\n",
    "    x = inputs\n",
    "    if conv_first:\n",
    "        x = conv(x)\n",
    "        if batch_normalization:\n",
    "            x = BatchNormalization()(x)\n",
    "        if activation is not None:\n",
    "            x = Activation(activation)(x)\n",
    "    else:\n",
    "        if batch_normalization:\n",
    "            x = BatchNormalization()(x)\n",
    "        else:\n",
    "            x = InstanceNormalization(x)\n",
    "        if activation is not None:\n",
    "            x = Activation(activation)(x)\n",
    "        x = conv(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RESNET V2 ARCHITECTURE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ResNet V2 architecture\n",
    "def resnet_v2(input_shape, depth, num_classes = 10):\n",
    "    if (depth - 2) % 9 != 0:\n",
    "        raise ValueError('Argument Error: \\'depth\\' argument should be 6n + 2 (eg 20, 32, 44 in [a])')\n",
    "    # Start model definition.\n",
    "    num_filters_in = 16\n",
    "    num_res_blocks = int((depth - 2) / 9)\n",
    "\n",
    "    inputs = Input(shape = input_shape)\n",
    "    # v2 performs Conv2D with BN-ReLU on input before splitting into 2 paths\n",
    "    x = FastConvolutionLayer(inputs = inputs,\n",
    "                    num_filters = num_filters_in,\n",
    "                    conv_first = True)\n",
    "\n",
    "    # Instantiate the stack of residual units\n",
    "    for stage in range(3):\n",
    "        for res_block in range(num_res_blocks):\n",
    "            activation = 'relu'\n",
    "            batch_normalization = True\n",
    "            strides = 1\n",
    "            if stage == 0:\n",
    "                num_filters_out = num_filters_in * 4\n",
    "                if res_block == 0: # first layer and first stage\n",
    "                    activation = None\n",
    "                    batch_normalization = False\n",
    "            else:\n",
    "                num_filters_out = num_filters_in * 2\n",
    "                if res_block == 0: # first layer but not first stage\n",
    "                    strides = 2 # downsample\n",
    "\n",
    "            # bottleneck residual unit\n",
    "            y = FastConvolutionLayer(inputs = x,\n",
    "                            num_filters = num_filters_in,\n",
    "                            kernel_size = 1,\n",
    "                            strides = strides,\n",
    "                            activation = activation,\n",
    "                            batch_normalization = batch_normalization,\n",
    "                            conv_first = False)\n",
    "            y = FastConvolutionLayer(inputs = y,\n",
    "                            num_filters = num_filters_in,\n",
    "                            conv_first = False)\n",
    "            y = FastConvolutionLayer(inputs = y,\n",
    "                            num_filters = num_filters_out,\n",
    "                            kernel_size = 1,\n",
    "                            conv_first = False)\n",
    "            if res_block == 0:\n",
    "                # Linear projection residual shortcut connection to match\n",
    "                # Changed dims\n",
    "                x = FastConvolutionLayer(inputs = x,\n",
    "                                        num_filters = num_filters_out,\n",
    "                                        kernel_size = 1,\n",
    "                                        strides = strides,\n",
    "                                        activation = None,\n",
    "                                        batch_normalization = False)\n",
    "            x = x + y\n",
    "\n",
    "        num_filters_in = num_filters_out\n",
    "\n",
    "    # Add classifier on top.\n",
    "    # v2 has BN-ReLU before Pooling\n",
    "    x = BatchNormalization()(x)\n",
    "    x = Activation('relu')(x)\n",
    "    x = AveragePooling2D(pool_size = 8)(x)\n",
    "    y = Flatten()(x)\n",
    "    outputs = Dense(num_classes,\n",
    "                    activation ='softmax',\n",
    "                    kernel_initializer ='he_normal')(y)\n",
    "\n",
    "    # Instantiate model.\n",
    "    model = Model(inputs = inputs, outputs = outputs)\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RESNET V1 ARCHITECTURE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ResNet V1 architecture\n",
    "class ResNet():\n",
    "def resnet_v1(input_shape, depth, num_classes = 10):\n",
    "    \n",
    "    if (depth - 2) % 6 != 0:\n",
    "        raise ValueError('Argument Error: \\'depth\\' argument should be 6n + 2 (eg 20, 32, 44 in [a])')\n",
    "    # Start model definition.\n",
    "    num_filters = 16\n",
    "    num_res_blocks = int((depth - 2) / 6)\n",
    "\n",
    "    inputs = Input(shape = input_shape)\n",
    "    x = resnet_layer(inputs = inputs)\n",
    "    # Instantiate the stack of residual units\n",
    "    for stack in range(3):\n",
    "        for res_block in range(num_res_blocks):\n",
    "            strides = 1\n",
    "            if stack > 0 and res_block == 0: # first layer but not first stack\n",
    "                strides = 2 # downsample\n",
    "            y = resnet_layer(inputs = x,\n",
    "                            num_filters = num_filters,\n",
    "                            strides = strides)\n",
    "            y = resnet_layer(inputs = y,\n",
    "                            num_filters = num_filters,\n",
    "                            activation = None)\n",
    "            if stack > 0 and res_block == 0: # first layer but not first stack\n",
    "                # linear projection residual shortcut connection to match\n",
    "                # changed dims\n",
    "                x = resnet_layer(inputs = x,\n",
    "                                num_filters = num_filters,\n",
    "                                kernel_size = 1,\n",
    "                                strides = strides,\n",
    "                                activation = None,\n",
    "                                batch_normalization = False)\n",
    "            x = keras.layers.add([x, y])\n",
    "            x = Activation('relu')(x)\n",
    "        num_filters *= 2\n",
    "\n",
    "    # Add classifier on top.\n",
    "    # v1 does not use BN after last shortcut connection-ReLU\n",
    "    x = AveragePooling2D(pool_size = 8)(x)\n",
    "    y = Flatten()(x)\n",
    "    outputs = Dense(num_classes,\n",
    "                    activation ='softmax',\n",
    "                    kernel_initializer ='he_normal')(y)\n",
    "\n",
    "    # Instantiate model.\n",
    "    model = Model(inputs = inputs, outputs = outputs)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GAN DEFINITION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GAN():\n",
    "    def __init__(self):\n",
    "        self.img_rows = 28\n",
    "        self.img_cols = 28\n",
    "        self.channels = 1\n",
    "        self.img_shape = (self.img_rows, self.img_cols, self.channels)\n",
    "        self.latent_dim = 100\n",
    "\n",
    "        optimizer = Adam(0.0002, 0.5)\n",
    "\n",
    "        # Build and compile the discriminator\n",
    "        self.discriminator = self.build_discriminator()\n",
    "        self.discriminator.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])\n",
    "\n",
    "        # Build the generator\n",
    "        self.generator = self.build_generator()\n",
    "\n",
    "        # The generator takes noise as input and generates imgs\n",
    "        z = Input(shape=(self.latent_dim,))\n",
    "        img = self.generator(z)\n",
    "\n",
    "        # For the combined model we will only train the generator\n",
    "        self.discriminator.trainable = False\n",
    "\n",
    "        # The discriminator takes generated images as input and determines validity\n",
    "        validity = self.discriminator(img)\n",
    "\n",
    "        # The combined model  (stacked generator and discriminator)\n",
    "        # Trains the generator to fool the discriminator\n",
    "        self.combined = Model(z, validity)\n",
    "        self.combined.compile(loss='binary_crossentropy', optimizer=optimizer)\n",
    "    \n",
    "\n",
    "\n",
    "'''     \n",
    "        Entradas = Input(shape=(filas,columnas,canales))\n",
    "        x = lrelu(fft_forward(Entradas))\n",
    "        x = lrelu(fft_forward(x))\n",
    "        x = SpectralPool2D(x)\n",
    "        x = instance_norm(x)\n",
    "        \n",
    "   '''     \n",
    "        \n",
    "        \n",
    "        \n",
    "def fft_model(img, weights, hparameters1, hparameters2):\n",
    "    fft = relu(fft_forward(img, weights[0], weights[1], hparameters1))\n",
    "    fft1 = relu(fft_forward(fft, weights[0], weights[1], hparameters1))\n",
    "    \n",
    "    \n",
    "    ##pool = pool_forward(fft.reshape(1,223,223,1), hparameters2)[0]\n",
    "    fft1 = relu(fft_forward(pool, weights[2], weights[3], hparameters1))\n",
    "    ##pool1 = pool_forward(fft1.reshape(1,fft1.shape[1],fft1.shape[1],1), hparameters2)[0]\n",
    "    fft2 = relu(fft_forward(pool1, weights[4], weights[5], hparameters1))\n",
    "    ##pool2 = pool_forward(fft2.reshape(1,fft2.shape[1],fft2.shape[1],1), hparameters2)[0]\n",
    "    ##flatten = pool2.reshape(1,pool2.shape[1]*pool2.shape[2]*pool2.shape[3])\n",
    "    ##neural_net = relu(np.dot(flatten, weights[6]) + weights[7])\n",
    "    ##ans = sigmoid(np.dot(neural_net, weights[8]) + weights[9])\n",
    "    \n",
    "    tup = (ans, fft, pool, fft1, pool1, fft2, pool2)\n",
    "    return tup\n",
    "        \n",
    "        \n",
    "        \n",
    "        model.add(Dense(512))\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        model.add(BatchNormalization(momentum=0.8))\n",
    "        model.add(Dense(1024))\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        model.add(BatchNormalization(momentum=0.8))\n",
    "        model.add(Dense(np.prod(self.img_shape), activation='tanh'))\n",
    "        model.add(Reshape(self.img_shape))\n",
    "\n",
    "        model.summary()\n",
    "\n",
    "        noise = Input(shape=(self.latent_dim,))\n",
    "        img = model(noise)\n",
    "\n",
    "        return Model(noise, img)\n",
    "\n",
    "        \n",
    "\n",
    "    def build_discriminator(self):\n",
    "\n",
    "        model = Sequential()\n",
    "\n",
    "        model.add(Flatten(input_shape=self.img_shape))\n",
    "        model.add(Dense(512))\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        model.add(Dense(256))\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        model.add(Dense(1, activation='sigmoid'))\n",
    "        model.summary()\n",
    "\n",
    "        img = Input(shape=self.img_shape)\n",
    "        validity = model(img)\n",
    "\n",
    "        return Model(img, validity)\n",
    "\n",
    "\n",
    "\n",
    "    def train(self, epochs, batch_size=128, sample_interval=50):\n",
    "        # Load the dataset\n",
    "        (X_train, _), (_, _) = mnist.load_data()\n",
    "\n",
    "        # Rescale -1 to 1\n",
    "        X_train = X_train / 127.5 - 1.\n",
    "        X_train = np.expand_dims(X_train, axis=3)\n",
    "        \n",
    "\n",
    "        # Adversarial ground truths\n",
    "        valid = np.ones((batch_size, 1))\n",
    "        fake = np.zeros((batch_size, 1))\n",
    "\n",
    "        for epoch in range(epochs):\n",
    "\n",
    "            # ---------------------\n",
    "            #  Train Discriminator\n",
    "            # ---------------------\n",
    "\n",
    "            # Select a random batch of images\n",
    "            idx = np.random.randint(0, X_train.shape[0], batch_size)\n",
    "            imgs = X_train[idx]\n",
    "\n",
    "            noise = np.random.normal(0, 1, (batch_size, self.latent_dim))\n",
    "\n",
    "            # Generate a batch of new images\n",
    "            gen_imgs = self.generator.predict(noise)\n",
    "\n",
    "            # Train the discriminator\n",
    "            d_loss_real = self.discriminator.train_on_batch(imgs, valid)\n",
    "            d_loss_fake = self.discriminator.train_on_batch(gen_imgs, fake)\n",
    "            d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)\n",
    "\n",
    "            # ---------------------\n",
    "            #  Train Generator\n",
    "            # ---------------------\n",
    "\n",
    "            noise = np.random.normal(0, 1, (batch_size, self.latent_dim))\n",
    "\n",
    "            # Train the generator (to have the discriminator label samples as valid)\n",
    "            g_loss = self.combined.train_on_batch(noise, valid)\n",
    "\n",
    "            # Plot the progress\n",
    "            print (\"%d [D loss: %f, acc.: %.2f%%] [G loss: %f]\" % (epoch, d_loss[0], 100*d_loss[1], g_loss))\n",
    "\n",
    "            # If at save interval => save generated image samples\n",
    "            if epoch % sample_interval == 0:\n",
    "                self.sample_images(epoch)\n",
    "\n",
    "\n",
    "\n",
    "    def sample_images(self, epoch):\n",
    "        r, c = 5, 5\n",
    "        noise = np.random.normal(0, 1, (r * c, self.latent_dim))\n",
    "        gen_imgs = self.generator.predict(noise)\n",
    "\n",
    "        # Rescale images 0 - 1\n",
    "        gen_imgs = 0.5 * gen_imgs + 0.5\n",
    "\n",
    "        fig, axs = plt.subplots(r, c)\n",
    "        cnt = 0\n",
    "        for i in range(r):\n",
    "            for j in range(c):\n",
    "                axs[i,j].imshow(gen_imgs[cnt, :,:,0], cmap='gray')\n",
    "                axs[i,j].axis('off')\n",
    "                cnt += 1\n",
    "        \n",
    "        im = '/Users/aleja/Downloads/gan_imgs/'\n",
    "        fig.savefig(im + \"%d.png\" % epoch)\n",
    "        plt.close()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    gan = GAN()\n",
    "    gan.train(epochs=1000, batch_size=32, sample_interval=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DECONVOLUTION LAYER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, unicode_literals, absolute_import, division\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from keras import backend as K\n",
    "from keras.engine.topology import Layer\n",
    "from scipy.fftpack import idct\n",
    "\n",
    "\n",
    "# create DCT filter basis as numpy matrix, one filter per column\n",
    "def dct_filters(filter_size=(3,3)):\n",
    "    N = filter_size[0] * filter_size[1]\n",
    "    filters = np.zeros((N,N-1), np.float32)\n",
    "    for i in range(1,N):\n",
    "        d = np.zeros(filter_size, np.float32)\n",
    "        d.flat[i] = 1\n",
    "        filters[:, i-1] = idct(idct(d, norm='ortho').T, norm='ortho').real.flatten()\n",
    "    return filters\n",
    "\n",
    "\n",
    "# Tensorflow: convert PSFs to OTFs\n",
    "# PSF: tensor with shape [height, width, channels_in, channels_out]\n",
    "# img_shape: pair of integers\n",
    "def psf2otf(psf, img_shape):\n",
    "    # Shape and type of the point spread function(s)\n",
    "    psf_shape = tf.shape(psf)\n",
    "    psf_type = psf.dtype\n",
    "\n",
    "    # Coordinates for 'cutting up' the PSF tensor\n",
    "    midH = tf.floor_div(psf_shape[0], 2)\n",
    "    midW = tf.floor_div(psf_shape[1], 2)\n",
    "\n",
    "    # Slice the PSF tensor into four parts\n",
    "    top_left     = psf[:midH, :midW, :, :]\n",
    "    top_right    = psf[:midH, midW:, :, :]\n",
    "    bottom_left  = psf[midH:, :midW, :, :]\n",
    "    bottom_right = psf[midH:, midW:, :, :]\n",
    "\n",
    "    # prepare zeros for filler\n",
    "    zeros_bottom = tf.zeros([psf_shape[0] - midH, img_shape[1] - psf_shape[1], psf_shape[2], psf_shape[3]], dtype=psf_type)\n",
    "    zeros_top    = tf.zeros([midH, img_shape[1] - psf_shape[1], psf_shape[2], psf_shape[3]], dtype=psf_type)\n",
    "\n",
    "    # construct top and bottom row of new tensor\n",
    "    top    = tf.concat([bottom_right, zeros_bottom, bottom_left], 1)\n",
    "    bottom = tf.concat([top_right,    zeros_top,    top_left],    1)\n",
    "\n",
    "    # prepare additional filler zeros and put everything together\n",
    "    zeros_mid = tf.zeros([img_shape[0] - psf_shape[0], img_shape[1], psf_shape[2], psf_shape[3]], dtype=psf_type)\n",
    "    pre_otf = tf.concat([top, zeros_mid, bottom], 0)\n",
    "    # output shape: [img_shape[0], img_shape[1], channels_in, channels_out]\n",
    "\n",
    "    # fast fourier transform, transposed because tensor must have shape [..., height, width] for this\n",
    "    otf = tf.fft2d(tf.cast(tf.transpose(pre_otf, perm=[2,3,0,1]), tf.complex64))\n",
    "\n",
    "    # output shape: [channels_in, channels_out, img_shape[0], img_shape[1]]\n",
    "    return otf\n",
    "\n",
    "\n",
    "class Pad(Layer):\n",
    "    def __init__(self, border=0, mode='REPLICATE', **kwargs):\n",
    "        assert border >= 0\n",
    "        assert mode in ['REPLICATE', 'ZEROS']\n",
    "        self.border = border\n",
    "        self.mode = mode\n",
    "        super(Pad, self).__init__(**kwargs)\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return (input_shape[0],\n",
    "                input_shape[1] + 2*self.border if input_shape[1] is not None else None,\n",
    "                input_shape[2] + 2*self.border if input_shape[2] is not None else None,\n",
    "                input_shape[3])\n",
    "\n",
    "    def call(self, x, mask=None):\n",
    "        if self.mode == 'REPLICATE':\n",
    "            # hack: iterate 1-pixel symmetric padding to get replicate padding\n",
    "            for i in range(self.border):\n",
    "                x = tf.pad(x, [[0,0], [1,1], [1,1], [0,0]], \"SYMMETRIC\")\n",
    "        elif self.mode == 'ZEROS':\n",
    "            x = tf.pad(x, [[0,0], [self.border, self.border], [self.border, self.border], [0,0]], \"CONSTANT\")\n",
    "        return x\n",
    "\n",
    "\n",
    "class Crop(Layer):\n",
    "    def __init__(self, border=0, **kwargs):\n",
    "        assert border >= 0\n",
    "        self.border = border\n",
    "        super(Crop, self).__init__(**kwargs)\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return (input_shape[0],\n",
    "                input_shape[1] - 2*self.border if input_shape[1] is not None else None,\n",
    "                input_shape[2] - 2*self.border if input_shape[2] is not None else None,\n",
    "                input_shape[3])\n",
    "\n",
    "    def call(self, x, mask=None):\n",
    "        return x[:, self.border:-self.border, self.border:-self.border, :] if self.border > 0 else x\n",
    "\n",
    "\n",
    "class FourierDeconvolution(Layer):\n",
    "    def __init__(self, filter_size, stage, **kwargs):\n",
    "        self.filter_size = filter_size\n",
    "        self.stage = stage\n",
    "        super(FourierDeconvolution, self).__init__(**kwargs)\n",
    "\n",
    "    def build(self, input_shapes):\n",
    "        # construct filter basis B and define filter weights variable\n",
    "        B = dct_filters(self.filter_size)\n",
    "        self.B = K.variable(B, name='B', dtype='float32')\n",
    "        self.nb_filters = B.shape[1]\n",
    "        self.filter_weights = K.variable(np.eye(self.nb_filters), name='filter_weights', dtype='float32')\n",
    "        self.trainable_weights = [self.filter_weights]\n",
    "        super(FourierDeconvolution, self).build(input_shapes)\n",
    "\n",
    "    def compute_output_shape(self, input_shapes):\n",
    "        return input_shapes[0]\n",
    "\n",
    "    def call(self, inputs, mask=None):\n",
    "        padded_inputs, adjustments, observations, blur_kernels, lambdas = inputs\n",
    "\n",
    "        imagesize = tf.shape(padded_inputs)[1:3]\n",
    "        kernelsize = tf.shape(blur_kernels)[1:3]\n",
    "        padding = tf.floor_div(kernelsize,2)\n",
    "\n",
    "        mask_int = tf.ones((imagesize[0]-2*padding[0],imagesize[1]-2*padding[1]), dtype=tf.float32)\n",
    "        mask_int = tf.pad(mask_int, [[padding[0],padding[0]],[padding[1],padding[1]]], mode='CONSTANT')\n",
    "        mask_int = tf.expand_dims(mask_int, 0)\n",
    "\n",
    "        filters = tf.matmul(self.B, self.filter_weights)\n",
    "        filters = tf.reshape(filters, [self.filter_size[0],self.filter_size[1],1,self.nb_filters])\n",
    "\n",
    "        filter_otfs = psf2otf(filters, imagesize)\n",
    "        otf_term = tf.reduce_sum(tf.square(tf.abs(filter_otfs)), axis=1)\n",
    "\n",
    "        k = tf.expand_dims(tf.transpose(blur_kernels, [1,2,0]), -1)\n",
    "        k_otf = psf2otf(k, imagesize)[:,0,:,:]\n",
    "\n",
    "        if self.stage > 1:\n",
    "            # boundary adjustment\n",
    "            Kx_fft       = tf.fft2d(tf.cast(padded_inputs[:,:,:,0], tf.complex64)) * k_otf\n",
    "            Kx           = tf.to_float(tf.ifft2d(Kx_fft))\n",
    "            Kx_outer     = (1.0 - mask_int) * Kx\n",
    "            y_inner      = mask_int * observations[:,:,:,0]\n",
    "            y_adjusted   = y_inner + Kx_outer\n",
    "            dataterm_fft = tf.fft2d(tf.cast(y_adjusted, tf.complex64)) * tf.conj(k_otf)\n",
    "        else:\n",
    "            # standard data term\n",
    "            observations_fft = tf.fft2d(tf.cast(observations[:,:,:,0], tf.complex64))\n",
    "            dataterm_fft = observations_fft * tf.conj(k_otf)\n",
    "\n",
    "        lambdas = tf.expand_dims(lambdas, -1)\n",
    "\n",
    "        adjustment_fft = tf.fft2d(tf.cast(adjustments[:,:,:,0], tf.complex64))\n",
    "        numerator_fft  = tf.cast(lambdas, tf.complex64) * dataterm_fft + adjustment_fft\n",
    "\n",
    "        KtK = tf.square(tf.abs(k_otf))\n",
    "        denominator_fft = lambdas * KtK + otf_term\n",
    "        denominator_fft = tf.cast(denominator_fft, tf.complex64)\n",
    "\n",
    "        frac_fft = numerator_fft / denominator_fft\n",
    "        return tf.expand_dims(tf.to_float(tf.ifft2d(frac_fft)), -1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}