r2u-net.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _recurrent_conv_layer(x,filters,kernel_size=3,padding=\"SAME\",strides=1,time_steps=2):\n",
    "    conv_base = layers.Conv2D(filters=filters,kernel_size=kernel_size,padding=padding)\n",
    "    x = conv_base(x)\n",
    "    norm = layers.BatchNormalization()\n",
    "    x = norm(x)\n",
    "    conv_t = layers.Conv2D(filters=filters,kernel_size=kernel_size,padding=padding)\n",
    "    \n",
    "    add = layers.Add()\n",
    "    x_t = tf.nn.relu(conv_t(x))#t0\n",
    "    x_t = norm(x_t)\n",
    "    for step in range(time_steps):\n",
    "        x_t = norm(tf.nn.relu(add([x,conv_t(x_t)])))\n",
    "    \n",
    "    return x_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def residual_recurrent_down_block(x,filters,kernel_size=3,padding=\"SAME\",strides=1,time_steps=2):\n",
    "    pool = layers.MaxPool2D((2,2),(2,2))\n",
    "    res = layers.Conv2D(filters=filters,kernel_size=1,padding=\"SAME\",strides=1,kernel_initializer=tf.constant_initializer())\n",
    "    add = layers.Add()\n",
    "    r = res(x)\n",
    "    x = _recurrent_conv_layer(x,filters,kernel_size,padding,strides,time_steps)\n",
    "    x = _recurrent_conv_layer(x,filters,kernel_size,padding,strides,time_steps)\n",
    "    \n",
    "    c = add([r,x])\n",
    "    p = pool(c)\n",
    "    return c,p\n",
    "\n",
    "def residual_recurrent_plateau_block(x,filters,kernel_size=3,padding=\"SAME\",strides=1,time_steps=2):\n",
    "    conv_transpose = layers.Conv2DTranspose(filters//2,kernel_size=(2,2),strides=(2,2))\n",
    "    res = layers.Conv2D(filters,kernel_size=1,padding=\"SAME\",strides=1,kernel_initializer=tf.constant_initializer())\n",
    "    add = layers.Add()\n",
    "    r = res(x)\n",
    "    x = _recurrent_conv_layer(x,filters,kernel_size,padding,strides,time_steps)\n",
    "    x = _recurrent_conv_layer(x,filters,kernel_size,padding,strides,time_steps)\n",
    "    x = add([r,x])\n",
    "    up = conv_transpose(x)\n",
    "    return up\n",
    "\n",
    "def residual_recurrent_up_block(x,p,filters,kernel_size=3,padding=\"SAME\",strides=1,time_steps=2):\n",
    "    res = layers.Conv2D(filters,kernel_size=1,padding=\"SAME\",strides=1,kernel_initializer=tf.constant_initializer())\n",
    "    add = layers.Add()\n",
    "    concat = layers.Concatenate()([x,p])\n",
    "    r = res(concat)\n",
    "    conv_transpose = layers.Conv2DTranspose(filters,kernel_size=(2,2),strides=(2,2))\n",
    "    x = _recurrent_conv_layer(concat,filters,kernel_size,padding,strides,time_steps)\n",
    "    x = _recurrent_conv_layer(x,filters,kernel_size,padding,strides,time_steps)\n",
    "    x = add([r,x])\n",
    "    up = conv_transpose(x)\n",
    "    return up\n",
    "\n",
    "def residual_recurrent_output_layer(x,p,filters,kernel_size=3,padding=\"SAME\",strides=1,time_steps=2):\n",
    "    concat = layers.Concatenate()([x,p])\n",
    "    out_conv = keras.layers.Conv2D(1, (1, 1), padding=\"same\", activation=\"relu\")\n",
    "    x = _recurrent_conv_layer(concat,filters,kernel_size,padding,strides,time_steps)\n",
    "    x = _recurrent_conv_layer(x,filters,kernel_size,padding,strides,time_steps)\n",
    "    outs = out_conv(x)\n",
    "    return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def R2UNet():\n",
    "    inputs = None\n",
    "    #downsampling\n",
    "    p0 = inputs\n",
    "    c1, p1 = residual_recurrent_down_block(p0, f[0]) #16\n",
    "    c2, p2 = residual_recurrent_down_block(p1, f[1]) #32\n",
    "    c3, p3 = residual_recurrent_down_block(p2, f[2]) #64\n",
    "    c4, p4 = residual_recurrent_down_block(p3, f[3]) #128\n",
    "    #convolution_block before upsampling\n",
    "    cn = residual_recurrent_plateau_block(p4, f[4]) #128\n",
    "    #upsampling\n",
    "    u1 = residual_recurrent_up_block(cn, c4, f[3]) #128\n",
    "    u2 = residual_recurrent_up_block(u1, c3, f[2]) #64\n",
    "    u3 = residual_recurrent_up_block(u2, c2, f[1]) #32\n",
    "    outputs = residual_recurrent_output_layer(u3, c1, f[0]) #16\n",
    "    \n",
    "    model = None#keras.models.Model(inputs, outputs)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}