unet_tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fdcdfd71",
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.models import Model\n",
    "from keras.layers import BatchNormalization, Input, Conv2D, MaxPooling2D, UpSampling2D, Activation, Concatenate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0c6a19c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv(x, filters, layername=None):\n",
    "    \"\"\"Convolutional layer with batch normalization and ReLU activation.\n",
    "    \n",
    "    :param filters: Number of filters\n",
    "    :return: output of the convolutional block\n",
    "    \"\"\"\n",
    "    x = Conv2D(filters=filters,\n",
    "               kernel_size=(3, 3),\n",
    "               padding='same',   \n",
    "               kernel_initializer='he_uniform',\n",
    "               name=layername)(x)\n",
    "    \n",
    "    x = BatchNormalization()(x)\n",
    "    x = Activation('relu')(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "429dfbfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Unet(filters=16, layers=4, input_shape=(None, None, 1), num_classes=1, activation='sigmoid'):\n",
    "    \"\"\"U-Net for semantic segmentation\n",
    "    \n",
    "    Reference:\n",
    "    Ronneberger, O., Fischer, P., & Brox, T. (2015, October).\n",
    "    U-net: Convolutional networks for biomedical image segmentation.\n",
    "    In International Conference on Medical image computing and computer-assisted\n",
    "    intervention (pp. 234-241). Springer, Cham\n",
    "    \n",
    "    :param filters: Number of filters in the first layer.\n",
    "    :param layers: Number of encoder and decoder layers\n",
    "    :param input_shape: Input shape (width, height, channels)\n",
    "    :param num_classes: Number of classes specified in last layer\n",
    "    :param activation: Activation of last layer\n",
    "    \n",
    "    :return: U-Net as Keras model\n",
    "    \"\"\"\n",
    "    \n",
    "    to_concat = []\n",
    "    \n",
    "    # Input layer\n",
    "    inputs = Input(input_shape, name='input_layer')\n",
    "    x = inputs\n",
    "    \n",
    "    # Encoder\n",
    "    for i in range(layers):\n",
    "        x = conv(x, filters * 2**i, f'enc_layer{i}_conv1')\n",
    "        x = conv(x, filters * 2**i, f'enc_layer{i}_conv2')\n",
    "        to_concat.append(x)\n",
    "        x = MaxPooling2D()(x)\n",
    "        \n",
    "    # Bottleneck\n",
    "    x = conv(x, filters * 2**(i+1), f'latent_conv')\n",
    "    \n",
    "    # Decoder\n",
    "    for i in range(layers):\n",
    "        x = UpSampling2D()(x)\n",
    "        x = Concatenate()([x, to_concat.pop()])\n",
    "        x = conv(x, filters * 2**(layers-1-i), f'dec_layer{i}_conv1')\n",
    "        x = conv(x, filters * 2**(layers-1-i), f'dec_layer{i}_conv2')\n",
    "    \n",
    "    # Output layer with 1x1 convolution\n",
    "    outputs = Conv2D(num_classes,\n",
    "                     kernel_size=(1, 1),\n",
    "                     padding='same',\n",
    "                     kernel_initializer='he_uniform',\n",
    "                     activation=activation,\n",
    "                     name='final_activation')(x)\n",
    "    \n",
    "    return Model(inputs, outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aab06b9d",
   "metadata": {},
   "source": [
    "**Unet_v2** below is less fancy, but maybe it is easier to understand (e.g. which layers' outputs are concatenated). <br> The architecture should be exactly the same as **Unet**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ee3abfcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Unet_v2(filters=16, layers=4, input_shape=(None, None, 1), num_classes=1, activation='sigmoid'):\n",
    "    \"\"\"U-Net for semantic segmentation\n",
    "    \n",
    "    Reference:\n",
    "    Ronneberger, O., Fischer, P., & Brox, T. (2015, October).\n",
    "    U-net: Convolutional networks for biomedical image segmentation.\n",
    "    In International Conference on Medical image computing and computer-assisted\n",
    "    intervention (pp. 234-241). Springer, Cham\n",
    "    \n",
    "    :param filters: Number of filters in the first layer.\n",
    "    :param layers: Number of encoder and decoder layers\n",
    "    :param input_shape: Input shape (width, height, channels)\n",
    "    :param num_classes: Number of classes specified in last layer\n",
    "    :param activation: Activation of last layer\n",
    "    \n",
    "    :return: U-Net as Keras model\n",
    "    \"\"\"\n",
    "    \n",
    "    # Input layer\n",
    "    inputs = Input(input_shape)\n",
    "    x = inputs\n",
    "    \n",
    "    # Encoder\n",
    "    c1 = conv(x, filters)\n",
    "    c1 = conv(c1, filters)\n",
    "    p1 = MaxPooling2D((2, 2))(c1)\n",
    "    \n",
    "    c2 = conv(p1, filters * 2)\n",
    "    c2 = conv(c2, filters * 2)\n",
    "    p2 = MaxPooling2D((2, 2))(c2)\n",
    "    \n",
    "    c3 = conv(p2, filters * 4)\n",
    "    c3 = conv(c3, filters * 4)\n",
    "    p3 = MaxPooling2D((2, 2))(c3)\n",
    "    \n",
    "    c4 = conv(p3, filters * 8)\n",
    "    c4 = conv(c4, filters * 8)\n",
    "    p4 = MaxPooling2D((2, 2))(c4)\n",
    "    \n",
    "    # Bottleneck\n",
    "    c5 = conv(p4, filters * 16)\n",
    "    \n",
    "    # Decoder\n",
    "    d1 = UpSampling2D()(c5)\n",
    "    d1 = Concatenate()([d1, c4])\n",
    "    d1 = conv(d1, filters * 8)\n",
    "    d1 = conv(d1, filters * 8)\n",
    "    \n",
    "    d2 = UpSampling2D()(d1)\n",
    "    d2 = Concatenate()([d2, c3])\n",
    "    d2 = conv(d2, filters * 4)\n",
    "    d2 = conv(d2, filters * 4)\n",
    "    \n",
    "    d3 = UpSampling2D()(d2)\n",
    "    d3 = Concatenate()([d3, c2])\n",
    "    d3 = conv(d3, filters * 2)\n",
    "    d3 = conv(d3, filters * 2)\n",
    "    \n",
    "    d4 = UpSampling2D()(d3)\n",
    "    d4 = Concatenate()([d4, c1])\n",
    "    d4 = conv(d4, filters)\n",
    "    d4 = conv(d4, filters)\n",
    "    \n",
    "    # Output layer with 1x1 convolution\n",
    "    outputs = Conv2D(num_classes,\n",
    "                     (1, 1),\n",
    "                     activation=activation,\n",
    "                     padding='same',\n",
    "                     kernel_initializer='he_uniform')(d4)\n",
    "    \n",
    "    return Model(inputs, outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3da5d41c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 1377121\n"
     ]
    }
   ],
   "source": [
    "# Example: Unet with 4 layers \n",
    "# for input images of size 32 x 32 with one color channel (grayscale)\n",
    "\n",
    "unet = Unet(filters=16, layers=4, input_shape=(32, 32, 1))\n",
    "\n",
    "print('Number of parameters:', unet.count_params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bc772f48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_layer (InputLayer)       [(None, 32, 32, 1)]  0           []                               \n",
      "                                                                                                  \n",
      " enc_layer0_conv1 (Conv2D)      (None, 32, 32, 16)   160         ['input_layer[0][0]']            \n",
      "                                                                                                  \n",
      " batch_normalization (BatchNorm  (None, 32, 32, 16)  64          ['enc_layer0_conv1[0][0]']       \n",
      " alization)                                                                                       \n",
      "                                                                                                  \n",
      " activation (Activation)        (None, 32, 32, 16)   0           ['batch_normalization[0][0]']    \n",
      "                                                                                                  \n",
      " enc_layer0_conv2 (Conv2D)      (None, 32, 32, 16)   2320        ['activation[0][0]']             \n",
      "                                                                                                  \n",
      " batch_normalization_1 (BatchNo  (None, 32, 32, 16)  64          ['enc_layer0_conv2[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_1 (Activation)      (None, 32, 32, 16)   0           ['batch_normalization_1[0][0]']  \n",
      "                                                                                                  \n",
      " max_pooling2d (MaxPooling2D)   (None, 16, 16, 16)   0           ['activation_1[0][0]']           \n",
      "                                                                                                  \n",
      " enc_layer1_conv1 (Conv2D)      (None, 16, 16, 32)   4640        ['max_pooling2d[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_2 (BatchNo  (None, 16, 16, 32)  128         ['enc_layer1_conv1[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_2 (Activation)      (None, 16, 16, 32)   0           ['batch_normalization_2[0][0]']  \n",
      "                                                                                                  \n",
      " enc_layer1_conv2 (Conv2D)      (None, 16, 16, 32)   9248        ['activation_2[0][0]']           \n",
      "                                                                                                  \n",
      " batch_normalization_3 (BatchNo  (None, 16, 16, 32)  128         ['enc_layer1_conv2[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_3 (Activation)      (None, 16, 16, 32)   0           ['batch_normalization_3[0][0]']  \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPooling2D)  (None, 8, 8, 32)    0           ['activation_3[0][0]']           \n",
      "                                                                                                  \n",
      " enc_layer2_conv1 (Conv2D)      (None, 8, 8, 64)     18496       ['max_pooling2d_1[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_4 (BatchNo  (None, 8, 8, 64)    256         ['enc_layer2_conv1[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_4 (Activation)      (None, 8, 8, 64)     0           ['batch_normalization_4[0][0]']  \n",
      "                                                                                                  \n",
      " enc_layer2_conv2 (Conv2D)      (None, 8, 8, 64)     36928       ['activation_4[0][0]']           \n",
      "                                                                                                  \n",
      " batch_normalization_5 (BatchNo  (None, 8, 8, 64)    256         ['enc_layer2_conv2[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_5 (Activation)      (None, 8, 8, 64)     0           ['batch_normalization_5[0][0]']  \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPooling2D)  (None, 4, 4, 64)    0           ['activation_5[0][0]']           \n",
      "                                                                                                  \n",
      " enc_layer3_conv1 (Conv2D)      (None, 4, 4, 128)    73856       ['max_pooling2d_2[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_6 (BatchNo  (None, 4, 4, 128)   512         ['enc_layer3_conv1[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_6 (Activation)      (None, 4, 4, 128)    0           ['batch_normalization_6[0][0]']  \n",
      "                                                                                                  \n",
      " enc_layer3_conv2 (Conv2D)      (None, 4, 4, 128)    147584      ['activation_6[0][0]']           \n",
      "                                                                                                  \n",
      " batch_normalization_7 (BatchNo  (None, 4, 4, 128)   512         ['enc_layer3_conv2[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_7 (Activation)      (None, 4, 4, 128)    0           ['batch_normalization_7[0][0]']  \n",
      "                                                                                                  \n",
      " max_pooling2d_3 (MaxPooling2D)  (None, 2, 2, 128)   0           ['activation_7[0][0]']           \n",
      "                                                                                                  \n",
      " latent_conv (Conv2D)           (None, 2, 2, 256)    295168      ['max_pooling2d_3[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_8 (BatchNo  (None, 2, 2, 256)   1024        ['latent_conv[0][0]']            \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_8 (Activation)      (None, 2, 2, 256)    0           ['batch_normalization_8[0][0]']  \n",
      "                                                                                                  \n",
      " up_sampling2d (UpSampling2D)   (None, 4, 4, 256)    0           ['activation_8[0][0]']           \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, 4, 4, 384)    0           ['up_sampling2d[0][0]',          \n",
      "                                                                  'activation_7[0][0]']           \n",
      "                                                                                                  \n",
      " dec_layer0_conv1 (Conv2D)      (None, 4, 4, 128)    442496      ['concatenate[0][0]']            \n",
      "                                                                                                  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " batch_normalization_9 (BatchNo  (None, 4, 4, 128)   512         ['dec_layer0_conv1[0][0]']       \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " activation_9 (Activation)      (None, 4, 4, 128)    0           ['batch_normalization_9[0][0]']  \n",
      "                                                                                                  \n",
      " dec_layer0_conv2 (Conv2D)      (None, 4, 4, 128)    147584      ['activation_9[0][0]']           \n",
      "                                                                                                  \n",
      " batch_normalization_10 (BatchN  (None, 4, 4, 128)   512         ['dec_layer0_conv2[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_10 (Activation)     (None, 4, 4, 128)    0           ['batch_normalization_10[0][0]'] \n",
      "                                                                                                  \n",
      " up_sampling2d_1 (UpSampling2D)  (None, 8, 8, 128)   0           ['activation_10[0][0]']          \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate)    (None, 8, 8, 192)    0           ['up_sampling2d_1[0][0]',        \n",
      "                                                                  'activation_5[0][0]']           \n",
      "                                                                                                  \n",
      " dec_layer1_conv1 (Conv2D)      (None, 8, 8, 64)     110656      ['concatenate_1[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_11 (BatchN  (None, 8, 8, 64)    256         ['dec_layer1_conv1[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_11 (Activation)     (None, 8, 8, 64)     0           ['batch_normalization_11[0][0]'] \n",
      "                                                                                                  \n",
      " dec_layer1_conv2 (Conv2D)      (None, 8, 8, 64)     36928       ['activation_11[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_12 (BatchN  (None, 8, 8, 64)    256         ['dec_layer1_conv2[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_12 (Activation)     (None, 8, 8, 64)     0           ['batch_normalization_12[0][0]'] \n",
      "                                                                                                  \n",
      " up_sampling2d_2 (UpSampling2D)  (None, 16, 16, 64)  0           ['activation_12[0][0]']          \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate)    (None, 16, 16, 96)   0           ['up_sampling2d_2[0][0]',        \n",
      "                                                                  'activation_3[0][0]']           \n",
      "                                                                                                  \n",
      " dec_layer2_conv1 (Conv2D)      (None, 16, 16, 32)   27680       ['concatenate_2[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_13 (BatchN  (None, 16, 16, 32)  128         ['dec_layer2_conv1[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_13 (Activation)     (None, 16, 16, 32)   0           ['batch_normalization_13[0][0]'] \n",
      "                                                                                                  \n",
      " dec_layer2_conv2 (Conv2D)      (None, 16, 16, 32)   9248        ['activation_13[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_14 (BatchN  (None, 16, 16, 32)  128         ['dec_layer2_conv2[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_14 (Activation)     (None, 16, 16, 32)   0           ['batch_normalization_14[0][0]'] \n",
      "                                                                                                  \n",
      " up_sampling2d_3 (UpSampling2D)  (None, 32, 32, 32)  0           ['activation_14[0][0]']          \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate)    (None, 32, 32, 48)   0           ['up_sampling2d_3[0][0]',        \n",
      "                                                                  'activation_1[0][0]']           \n",
      "                                                                                                  \n",
      " dec_layer3_conv1 (Conv2D)      (None, 32, 32, 16)   6928        ['concatenate_3[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_15 (BatchN  (None, 32, 32, 16)  64          ['dec_layer3_conv1[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_15 (Activation)     (None, 32, 32, 16)   0           ['batch_normalization_15[0][0]'] \n",
      "                                                                                                  \n",
      " dec_layer3_conv2 (Conv2D)      (None, 32, 32, 16)   2320        ['activation_15[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_16 (BatchN  (None, 32, 32, 16)  64          ['dec_layer3_conv2[0][0]']       \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " activation_16 (Activation)     (None, 32, 32, 16)   0           ['batch_normalization_16[0][0]'] \n",
      "                                                                                                  \n",
      " final_activation (Conv2D)      (None, 32, 32, 1)    17          ['activation_16[0][0]']          \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 1,377,121\n",
      "Trainable params: 1,374,689\n",
      "Non-trainable params: 2,432\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "unet.summary()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}