careful.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras.layers as layers\n",
    "import tensorflow.keras.models as models\n",
    "import tensorflow.keras.activations as activations\n",
    "import tensorflow.keras.metrics as metrics\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "import tensorflow_addons.optimizers as optimizers\n",
    "import tensorflow_addons.losses as losses\n",
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "\n",
    "# Loads dataset.\n",
    "# Targets are the segmantation images.\n",
    "# Inputs are the original images.\n",
    "def load_dataset(dataset_path):\n",
    "    data = np.load(dataset_path, allow_pickle=True)\n",
    "    return data['inputs'], data['targets']\n",
    "\n",
    "\n",
    "# Defining the encoder's down-sampling blocks.\n",
    "def encoder_block(inputs, n_filters, kernel_size, strides):\n",
    "    encoder = layers.Conv2D(filters=n_filters, kernel_size=kernel_size, strides=strides, padding='same', use_bias=False)(inputs)\n",
    "    encoder = layers.BatchNormalization()(encoder)\n",
    "    encoder = layers.Activation(activations.gelu)(encoder)\n",
    "    encoder = layers.Conv2D(filters=n_filters, kernel_size=kernel_size, padding='same', use_bias=False)(encoder)\n",
    "    encoder = layers.BatchNormalization()(encoder)\n",
    "    encoder = layers.Activation(activations.gelu)(encoder)\n",
    "    return encoder\n",
    "\n",
    "\n",
    "# 定义解码器的上采样块。\n",
    "def upscale_blocks(inputs):\n",
    "    n_upscales = len(inputs)\n",
    "    upscale_layers = []\n",
    "\n",
    "    for i, inp in enumerate(inputs):\n",
    "        p = n_upscales - i\n",
    "        u = layers.Conv2DTranspose(filters=64, kernel_size=3, strides=2**p, padding='same')(inp)\n",
    "\n",
    "        for i in range(2):\n",
    "            u = layers.Conv2D(filters=64, kernel_size=3, padding='same', use_bias=False)(u)\n",
    "            u = layers.BatchNormalization()(u)\n",
    "            u = layers.Activation(activations.gelu)(u)\n",
    "            u = layers.Dropout(rate=0.4)(u)\n",
    "\n",
    "        upscale_layers.append(u)\n",
    "    return upscale_layers\n",
    "\n",
    "\n",
    "# 定义解码器的整个模块。\n",
    "\n",
    "def decoder_block(layers_to_upscale, inputs):\n",
    "    upscaled_layers = upscale_blocks(layers_to_upscale)\n",
    "\n",
    "    decoder_blocks = []\n",
    "\n",
    "    for i, inp in enumerate(inputs):\n",
    "        d = layers.Conv2D(filters=64, kernel_size=3, strides=2**i, padding='same', use_bias=False)(inp)\n",
    "        d = layers.BatchNormalization()(d)\n",
    "        d = layers.Activation(activations.gelu)(d)\n",
    "        d = layers.Conv2D(filters=64, kernel_size=3, padding='same', use_bias=False)(d)\n",
    "        d = layers.BatchNormalization()(d)\n",
    "        d = layers.Activation(activations.gelu)(d)\n",
    "\n",
    "        decoder_blocks.append(d)\n",
    "\n",
    "    decoder = layers.concatenate(upscaled_layers + decoder_blocks)\n",
    "    decoder = layers.Conv2D(filters=256, kernel_size=3, strides=1, padding='same', use_bias=False)(decoder)\n",
    "    decoder = layers.BatchNormalization()(decoder)\n",
    "    decoder = layers.Activation(activations.gelu)(decoder)\n",
    "    decoder = layers.Dropout(rate=0.4)(decoder)\n",
    "\n",
    "    return decoder\n",
    "\n",
    "\n",
    "def get_model(input_dim):\n",
    "    inputs = layers.Input(input_dim)\n",
    "\n",
    "    noisy_inputs = layers.GaussianNoise(stddev=0.2)(inputs)\n",
    "\n",
    "    e1 = encoder_block(noisy_inputs, n_filters=32, kernel_size=3, strides=1)\n",
    "    e2 = encoder_block(e1, n_filters=64, kernel_size=3, strides=2)\n",
    "    e3 = encoder_block(e2, n_filters=128, kernel_size=3, strides=2)\n",
    "    e4 = encoder_block(e3, n_filters=256, kernel_size=3, strides=2)\n",
    "    e5 = encoder_block(e4, n_filters=512, kernel_size=3, strides=2)\n",
    "\n",
    "    d4 = decoder_block(layers_to_upscale=[e5], inputs=[e4, e3, e2, e1])\n",
    "    d3 = decoder_block(layers_to_upscale=[e5, d4], inputs=[e3, e2, e1])\n",
    "    d2 = decoder_block(layers_to_upscale=[e5, d4, d3], inputs=[e2, e1])\n",
    "    d1 = decoder_block(layers_to_upscale=[e5, d4, d3, d2], inputs=[e1])\n",
    "\n",
    "    output = layers.Conv2D(filters=3, kernel_size=1, padding='same', activation='tanh')(d1)\n",
    "\n",
    "    model = models.Model(inputs, output)\n",
    "    return model\n",
    "\n",
    "\n",
    "model = get_model((256, 256, 3))\n",
    "model.summary()\n",
    "\n",
    "\n",
    "# Loading the dataset.\n",
    "# x_train: Segmentation images.\n",
    "# y_train: Original images.\n",
    "DATASET_PATH = 'D:\\\\Datasets\\\\City-Segmentation\\\\Cityscapes\\\\dataset.npz'\n",
    "x_train, y_train = load_dataset(DATASET_PATH)\n",
    "\n",
    "# Normalizing data.\n",
    "x_train = (x_train - 127.5) / 127.5\n",
    "y_train = (y_train - 127.5) / 127.5\n",
    "\n",
    "# Building & Compiling the model.\n",
    "image_size = x_train[0].shape\n",
    "unet3_plus = get_model(image_size)\n",
    "\n",
    "unet3_plus.compile(\n",
    "    optimizer=optimizers.Yogi(learning_rate=0.00025),\n",
    "    loss=losses.sigmoid_focal_crossentropy,\n",
    "    metrics=[metrics.MeanIoU(num_classes=30)]\n",
    ")\n",
    "\n",
    "# Training the model\n",
    "epochs = 200\n",
    "batch_size = 4\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(buffer_size=x_train.shape[0])\n",
    "inputs = dataset.batch(batch_size=batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)\n",
    "batches_per_epoch = x_train.shape[0] // batch_size\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    print('\\nTraining on epoch', epoch + 1)\n",
    "\n",
    "    loss = 0\n",
    "    for i, (x_batch, y_batch) in enumerate(inputs):\n",
    "        loss = unet3_plus.train_on_batch(x_batch, y_batch)\n",
    "\n",
    "        print('\\rCurrent batch: {}/{} , loss = {}'.format(\n",
    "            i+1,\n",
    "            batches_per_epoch,\n",
    "            loss, end='')\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "  # def add_prediction_op(self):\n",
    "  #   logging.info(\"Model: depths {depths}, filters {filters}, \"\n",
    "  #          \"filter size {kernel_size[0]}x{kernel_size[1]}, \"\n",
    "  #          \"pool size: {pool_size[0]}x{pool_size[1]}, \"\n",
    "  #          \"dilation rate: {dilation_rate[0]}x{dilation_rate[1]}\".format(\n",
    "  #           depths=self.depths,\n",
    "  #           filters=self.filters_root,\n",
    "  #           kernel_size=self.kernel_size,\n",
    "  #           dilation_rate=self.dilation_rate,\n",
    "  #           pool_size=self.pool_size))\n",
    "\n",
    "\n",
    "  #   # 如果 weight_decay 大于 0，设置 L2 正则化以防止过拟合，否则不使用正则化\n",
    "  #   if self.weight_decay > 0:\n",
    "  #     weight_decay = tf.constant(self.weight_decay, dtype=tf.float32, name=\"weight_constant\")\n",
    "  #     self.regularizer = tf.keras.regularizers.l2(l=0.5 * (weight_decay))\n",
    "  #   else:\n",
    "  #     self.regularizer = None\n",
    "\n",
    "  #   #权重初始化器，1.0不缩放\n",
    "  #   self.initializer = tf.compat.v1.keras.initializers.VarianceScaling(scale=1.0, mode=\"fan_avg\", distribution=\"uniform\")\n",
    "\n",
    "\n",
    "\n",
    "  #   # 输入层第一部分（3*3001-->8*3001）\n",
    "  #   convs = [None] * self.depths # 存储每个深度的输出\n",
    "\n",
    "  #   with tf.compat.v1.variable_scope(\"Input\"):\n",
    "  #     net = self.X\n",
    "\n",
    "  #     #对输入数据应用卷积操作，滤波器数量为 filters_root，卷积核大小为 kernel_size，padding='same' 保持输出的空间尺寸与输入相同。\n",
    "  #     net = tf.compat.v1.layers.conv2d(net,\n",
    "  #                  filters=self.filters_root,\n",
    "  #                  kernel_size=self.kernel_size,\n",
    "  #                  activation=None,\n",
    "  #                  padding='same',\n",
    "  #                  dilation_rate=self.dilation_rate,\n",
    "  #                  kernel_initializer=self.initializer,\n",
    "  #                  kernel_regularizer=self.regularizer,\n",
    "  #                  name=\"input_conv\")\n",
    "      \n",
    "  #     #对卷积后的特征图进行批量归一化，以稳定训练过程，加快收敛。\n",
    "  #     net = tf.compat.v1.layers.batch_normalization(net,\n",
    "  #                       training=self.is_training,\n",
    "  #                       name=\"input_bn\")\n",
    "  #     #应用 ReLU 激活函数，将非线性引入网络，增加模型表达能力。\n",
    "  #     net = tf.nn.relu(net,\n",
    "  #              name=\"input_relu\")\n",
    "  #     # net = tf.nn.dropout(net, self.keep_prob)\n",
    "      \n",
    "  #     #对特征图应用 Dropout，以防止过拟合，drop_rate 控制丢弃的比例。\n",
    "  #     net = tf.compat.v1.layers.dropout(net,\n",
    "  #                 rate=self.drop_rate,\n",
    "  #                 training=self.is_training,\n",
    "  #                 name=\"input_dropout\")\n",
    "\n",
    "\n",
    "\n",
    "  #   #上面为输入层，现在开始下采样层（包含了输入层第二部分(8*3001-->8*3001)）\n",
    "  #   for depth in range(0, self.depths):\n",
    "  #     with tf.compat.v1.variable_scope(\"DownConv_%d\" % depth):\n",
    "        \n",
    "  #       filters = int(2**(depth) * self.filters_root)\n",
    "        \n",
    "  #       net = tf.compat.v1.layers.conv2d(net,\n",
    "  #                    filters=filters,\n",
    "  #                    kernel_size=self.kernel_size,\n",
    "  #                    activation=None,\n",
    "  #                    use_bias=False,\n",
    "  #                    padding='same',\n",
    "  #                    dilation_rate=self.dilation_rate,\n",
    "  #                    kernel_initializer=self.initializer,\n",
    "  #                    kernel_regularizer=self.regularizer,\n",
    "  #                    name=\"down_conv1_{}\".format(depth + 1))\n",
    "  #       net = tf.compat.v1.layers.batch_normalization(net,\n",
    "  #                         training=self.is_training,\n",
    "  #                         name=\"down_bn1_{}\".format(depth + 1))\n",
    "  #       net = tf.nn.relu(net,\n",
    "  #                name=\"down_relu1_{}\".format(depth+1))\n",
    "  #       net = tf.compat.v1.layers.dropout(net,\n",
    "  #                   rate=self.drop_rate,\n",
    "  #                   training=self.is_training,\n",
    "  #                   name=\"down_dropout1_{}\".format(depth + 1))\n",
    "\n",
    "  #       convs[depth] = net \n",
    "  #       #每层网络结构convs储存\n",
    "  #       #convs[0]是第二个8*3001,\n",
    "  #       #convs[1]是11*751\n",
    "  #       #convs[2]is 16*188\n",
    "  #       #convs[3]is 22*47\n",
    "  #       #convs[4]is 32*12\n",
    "  #       #此刻convs储存了所有的需要skip connection的块\n",
    "\n",
    "\n",
    "  #       #下层前四层0，1，2，3\n",
    "  #       if depth < self.depths - 1:\n",
    "  #         net = tf.compat.v1.layers.conv2d(net,\n",
    "  #                      filters=filters,\n",
    "  #                      kernel_size=self.kernel_size,\n",
    "  #                      strides=self.pool_size, #Convolution+Stride+Rule\n",
    "  #                      activation=None,\n",
    "  #                      use_bias=False,\n",
    "  #                      padding='same',\n",
    "  #                      dilation_rate=self.dilation_rate,\n",
    "  #                      kernel_initializer=self.initializer,\n",
    "  #                      kernel_regularizer=self.regularizer,\n",
    "  #                      name=\"down_conv3_{}\".format(depth + 1))\n",
    "  #         net = tf.compat.v1.layers.batch_normalization(net,\n",
    "  #                           training=self.is_training,\n",
    "  #                           name=\"down_bn3_{}\".format(depth + 1))\n",
    "  #         net = tf.nn.relu(net,\n",
    "  #                  name=\"down_relu3_{}\".format(depth+1))\n",
    "  #         net = tf.compat.v1.layers.dropout(net,\n",
    "  #                   rate=self.drop_rate,\n",
    "  #                   training=self.is_training,\n",
    "  #                   name=\"down_dropout3_{}\".format(depth + 1))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "  #   # 上层（4，3，2，1层）（3，2，1，0）\n",
    "  #   for depth in range(self.depths - 2, -1, -1):\n",
    "  #     with tf.compat.v1.variable_scope(\"UpConv_%d\" % depth):\n",
    "  #       filters = int(2**(depth) * self.filters_root) #64 32 16 8\n",
    "  #       #卷积核随着深度增加而变化\n",
    "  #       net = tf.compat.v1.layers.conv2d_transpose(net,\n",
    "  #                        filters=filters,\n",
    "  #                        kernel_size=self.kernel_size,\n",
    "  #                        strides=self.pool_size,\n",
    "  #                        activation=None,\n",
    "  #                        use_bias=False,\n",
    "  #                        padding=\"same\",\n",
    "  #                        kernel_initializer=self.initializer,\n",
    "  #                        kernel_regularizer=self.regularizer,\n",
    "  #                        name=\"up_conv0_{}\".format(depth+1))\n",
    "  #       net = tf.compat.v1.layers.batch_normalization(net,\n",
    "  #                         training=self.is_training,\n",
    "  #                         name=\"up_bn0_{}\".format(depth + 1))\n",
    "  #       net = tf.nn.relu(net,\n",
    "  #                name=\"up_relu0_{}\".format(depth+1))\n",
    "  #       net = tf.compat.v1.layers.dropout(net,\n",
    "  #                   rate=self.drop_rate,\n",
    "  #                   training=self.is_training,\n",
    "  #                   name=\"up_dropout0_{}\".format(depth + 1))\n",
    "\n",
    "        \n",
    "  #       #跳跃连接\n",
    "  #       #net>convs[depth]--size\n",
    "  #       net = crop_and_concat(convs[depth], net)\n",
    "  #       #net = crop_only(convs[depth], net)\n",
    "\n",
    "  #       net = tf.compat.v1.layers.conv2d(net,\n",
    "  #                    filters=filters,\n",
    "  #                    kernel_size=self.kernel_size,\n",
    "  #                    activation=None,\n",
    "  #                    use_bias=False,\n",
    "  #                    padding='same',\n",
    "  #                    dilation_rate=self.dilation_rate,\n",
    "  #                    kernel_initializer=self.initializer,\n",
    "  #                    kernel_regularizer=self.regularizer,\n",
    "  #                    name=\"up_conv1_{}\".format(depth + 1))\n",
    "  #       net = tf.compat.v1.layers.batch_normalization(net,\n",
    "  #                         training=self.is_training,\n",
    "  #                         name=\"up_bn1_{}\".format(depth + 1))\n",
    "  #       net = tf.nn.relu(net,\n",
    "  #                name=\"up_relu1_{}\".format(depth + 1))\n",
    "  #       net = tf.compat.v1.layers.dropout(net,\n",
    "  #                   rate=self.drop_rate,\n",
    "  #                   training=self.is_training,\n",
    "  #                   name=\"up_dropout1_{}\".format(depth + 1))\n",
    "\n",
    "\n",
    "  #   # Output Map\n",
    "  #   with tf.compat.v1.variable_scope(\"Output\"):\n",
    "  #     net = tf.compat.v1.layers.conv2d(net,\n",
    "  #                  filters=self.n_class,\n",
    "  #                  kernel_size=(1,1),\n",
    "  #                  activation=None,\n",
    "  #                  padding='same',\n",
    "  #                  #dilation_rate=self.dilation_rate,\n",
    "  #                  kernel_initializer=self.initializer,\n",
    "  #                  kernel_regularizer=self.regularizer,\n",
    "  #                  name=\"output_conv\")\n",
    "  #     # net = tf.nn.relu(net,\n",
    "  #     #                     name=\"output_relu\")\n",
    "  #     # net = tf.compat.v1.layers.dropout(net,\n",
    "  #     #                         rate=self.drop_rate,\n",
    "  #     #                         training=self.is_training,\n",
    "  #     #                         name=\"output_dropout\")\n",
    "  #     # net = tf.compat.v1.layers.batch_normalization(net,\n",
    "  #     #                                    training=self.is_training,\n",
    "  #     #                                    name=\"output_bn\")\n",
    "  #     output = net\n",
    "\n",
    "\n",
    "\n",
    "  #   # 保存编码器部分的最终表示，可以用于其他任务或分析\n",
    "  #   with tf.compat.v1.variable_scope(\"representation\"):\n",
    "  #     self.representation = convs[-1]\n",
    "\n",
    "  #   # logits 是最终的网络输出，未经过激活函数\n",
    "  #   with tf.compat.v1.variable_scope(\"logits\"):\n",
    "  #     self.logits = output\n",
    "  #     # 记录 logits 的直方图摘要，以便在 TensorBoard 中可视化\n",
    "  #     tmp = tf.compat.v1.summary.histogram(\"logits\", self.logits)\n",
    "  #     self.summary_train.append(tmp)\n",
    "\n",
    "  #   # 计算最终预测结果，使用 softmax 将 logits 转换为概率分布\n",
    "  #   with tf.compat.v1.variable_scope(\"preds\"):\n",
    "  #     self.preds = tf.nn.softmax(output)\n",
    "  #     # 记录预测结果的直方图摘要，以便在 TensorBoard 中可视化\n",
    "  #     tmp = tf.compat.v1.summary.histogram(\"preds\", self.preds)\n",
    "  #     self.summary_train.append(tmp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "  # def add_prediction_op(self):\n",
    "  #   logging.info(\"Model: depths {depths}, filters {filters}, \"\n",
    "  #          \"filter size {kernel_size[0]}x{kernel_size[1]}, \"\n",
    "  #          \"pool size: {pool_size[0]}x{pool_size[1]}, \"\n",
    "  #          \"dilation rate: {dilation_rate[0]}x{dilation_rate[1]}\".format(\n",
    "  #           depths=self.depths,\n",
    "  #           filters=self.filters_root,\n",
    "  #           kernel_size=self.kernel_size,\n",
    "  #           dilation_rate=self.dilation_rate,\n",
    "  #           pool_size=self.pool_size))\n",
    "\n",
    "\n",
    "  #   # 如果 weight_decay 大于 0，设置 L2 正则化以防止过拟合，否则不使用正则化\n",
    "  #   # 设置 L2 正则化\n",
    "  #   self.regularizer = tf.keras.regularizers.l2(0.5 * self.weight_decay) if self.weight_decay > 0 else None\n",
    "\n",
    "\n",
    "  #   # 初始化器\n",
    "  #   self.initializer = tf.keras.initializers.VarianceScaling(scale=1.0, mode=\"fan_avg\", distribution=\"uniform\")\n",
    "\n",
    "\n",
    "\n",
    "  #   # 输入层第一部分（3*3001-->8*3001）\n",
    "  #   convs = [None] * self.depths # 存储每个深度的输出\n",
    "\n",
    "  #   with tf.name_scope(\"Input\"):\n",
    "  #     net = self.X  # 如果 self.X 是输入张量，直接赋值给 net\n",
    "\n",
    "  #     net = tf.keras.layers.Conv2D(\n",
    "  #                                  filters=self.filters_root, kernel_size=self.kernel_size, padding='same', dilation_rate=self.dilation_rate, kernel_initializer=self.initializer,kernel_regularizer=self.regularizer,name=\"input_conv\")(net)\n",
    "  #     net = tf.keras.layers.BatchNormalization( name=\"input_bn\")(net, training=self.is_training)\n",
    "  #     net = tf.keras.layers.ReLU(\n",
    "  #                                name=\"input_relu\")(net)\n",
    "  #     net = tf.keras.layers.Dropout(rate=self.drop_rate)(net, training=self.is_training)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "  #   # 下采样部分\n",
    "  #   # 第一层\n",
    "  #   with tf.name_scope(\"DownConv_0\"):\n",
    "  #       filters = int(2**0 * self.filters_root)  # 8\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_0_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_0_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_0_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_0_1\")(net, training=self.is_training)\n",
    "\n",
    "  #       convs[0] = net  # 存储跳跃连接块\n",
    "\n",
    "  #       # 下采样卷积\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   strides=self.pool_size,  # Convolution+Stride\n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_0_2\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_0_2\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_0_2\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_0_2\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第二层\n",
    "  #   with tf.name_scope(\"DownConv_1\"):\n",
    "  #       filters = int(2**1 * self.filters_root)  # 16\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_1_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_1_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_1_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_1_1\")(net, training=self.is_training)\n",
    "  #       convs[1] = net  # 存储跳跃连接块\n",
    "\n",
    "  #       # 下采样卷积\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   strides=self.pool_size,  # Convolution+Stride\n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_1_2\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_1_2\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_1_2\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_1_2\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第三层\n",
    "  #   with tf.name_scope(\"DownConv_2\"):\n",
    "  #       filters = int(2**2 * self.filters_root)  # 32\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_2_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_2_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_2_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_2_1\")(net, training=self.is_training)\n",
    "  #       convs[2] = net  # 存储跳跃连接块\n",
    "\n",
    "  #       # 下采样卷积\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   strides=self.pool_size,  # Convolution+Stride\n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_2_2\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_2_2\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_2_2\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_2_2\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第四层\n",
    "  #   with tf.name_scope(\"DownConv_3\"):\n",
    "  #       filters = int(2**3 * self.filters_root)  # 64\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_3_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_3_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_3_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_3_1\")(net, training=self.is_training)\n",
    "  #       convs[3] = net  # 存储跳跃连接块\n",
    "\n",
    "  #       # 下采样卷积\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   strides=self.pool_size,  # Convolution+Stride\n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_3_2\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_3_2\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_3_2\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_3_2\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第五层\n",
    "  #   with tf.name_scope(\"DownConv_4\"):\n",
    "  #       filters = int(2**4 * self.filters_root)  # 128\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   use_bias=False, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"down_conv_4_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"down_bn_4_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"down_relu_4_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"down_dropout_4_1\")(net, training=self.is_training)\n",
    "  #       convs[4] = net  # 存储跳跃连接块\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "  #     ## 上采样和跳跃连接部分\n",
    "  #   # 第四层（从上至下，编号为 3）\n",
    "  #   with tf.name_scope(\"UpConv_3\"):\n",
    "  #       filters = int(2**3 * self.filters_root)  # 64\n",
    "  #       net = tf.keras.layers.Conv2DTranspose(filters=filters, \n",
    "  #                                             kernel_size=self.kernel_size, \n",
    "  #                                             strides=self.pool_size, \n",
    "  #                                             padding=\"same\", \n",
    "  #                                             use_bias=False, \n",
    "  #                                             kernel_initializer=self.initializer, \n",
    "  #                                             kernel_regularizer=self.regularizer, \n",
    "  #                                             name=\"up_conv0_3\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn0_3\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu0_3\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout0_3\")(net, training=self.is_training)\n",
    "  #       #convs--8 16 32 64 128\n",
    "  #       # 跳跃连接\n",
    "  #       net = crop_and_concat(convs[3], net)\n",
    "\n",
    "  #       # 卷积操作\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   use_bias=False, \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"up_conv1_3\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn1_3\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu1_3\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout1_3\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第三层\n",
    "  #   with tf.name_scope(\"UpConv_2\"):\n",
    "  #       filters = int(2**2 * self.filters_root)  # 32\n",
    "  #       net = tf.keras.layers.Conv2DTranspose(filters=filters, \n",
    "  #                                             kernel_size=self.kernel_size, \n",
    "  #                                             strides=self.pool_size, \n",
    "  #                                             padding=\"same\", \n",
    "  #                                             use_bias=False, \n",
    "  #                                             kernel_initializer=self.initializer, \n",
    "  #                                             kernel_regularizer=self.regularizer, \n",
    "  #                                             name=\"up_conv0_2\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn0_2\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu0_2\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout0_2\")(net, training=self.is_training)\n",
    "\n",
    "  #       # 跳跃连接\n",
    "  #       net = crop_and_concat(convs[2], net)\n",
    "\n",
    "  #       # 卷积操作\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   use_bias=False, \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"up_conv1_2\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn1_2\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu1_2\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout1_2\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第二层\n",
    "  #   with tf.name_scope(\"UpConv_1\"):\n",
    "  #       filters = int(2**1 * self.filters_root)  # 16\n",
    "  #       net = tf.keras.layers.Conv2DTranspose(filters=filters, \n",
    "  #                                             kernel_size=self.kernel_size, \n",
    "  #                                             strides=self.pool_size, \n",
    "  #                                             padding=\"same\", \n",
    "  #                                             use_bias=False, \n",
    "  #                                             kernel_initializer=self.initializer, \n",
    "  #                                             kernel_regularizer=self.regularizer, \n",
    "  #                                             name=\"up_conv0_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn0_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu0_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout0_1\")(net, training=self.is_training)\n",
    "\n",
    "  #       # 跳跃连接\n",
    "  #       net = crop_and_concat(convs[1], net)\n",
    "\n",
    "  #       # 卷积操作\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   use_bias=False, \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"up_conv1_1\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn1_1\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu1_1\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout1_1\")(net, training=self.is_training)\n",
    "\n",
    "  #   # 第一层\n",
    "  #   with tf.name_scope(\"UpConv_0\"):\n",
    "  #       filters = int(2**0 * self.filters_root)  # 8\n",
    "  #       net = tf.keras.layers.Conv2DTranspose(filters=filters, \n",
    "  #                                             kernel_size=self.kernel_size, \n",
    "  #                                             strides=self.pool_size, \n",
    "  #                                             padding=\"same\", \n",
    "  #                                             use_bias=False, \n",
    "  #                                             kernel_initializer=self.initializer, \n",
    "  #                                             kernel_regularizer=self.regularizer, \n",
    "  #                                             name=\"up_conv0_0\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn0_0\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu0_0\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout0_0\")(net, training=self.is_training)\n",
    "\n",
    "  #       # 跳跃连接\n",
    "  #       net = crop_and_concat(convs[0], net)\n",
    "\n",
    "  #       # 卷积操作\n",
    "  #       net = tf.keras.layers.Conv2D(filters=filters, \n",
    "  #                                   kernel_size=self.kernel_size, \n",
    "  #                                   padding='same', \n",
    "  #                                   use_bias=False, \n",
    "  #                                   dilation_rate=self.dilation_rate, \n",
    "  #                                   kernel_initializer=self.initializer, \n",
    "  #                                   kernel_regularizer=self.regularizer, \n",
    "  #                                   name=\"up_conv1_0\")(net)\n",
    "  #       net = tf.keras.layers.BatchNormalization(name=\"up_bn1_0\")(net, training=self.is_training)\n",
    "  #       net = tf.keras.layers.ReLU(name=\"up_relu1_0\")(net)\n",
    "  #       net = tf.keras.layers.Dropout(rate=self.drop_rate, name=\"up_dropout1_0\")(net, training=self.is_training)\n",
    "\n",
    "\n",
    "  #   #全连接\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "  #   # Output Map\n",
    "  #   # 输出层\n",
    "  #   with tf.name_scope(\"Output\"):\n",
    "  #       net = tf.keras.layers.Conv2D(filters=self.n_class,\n",
    "  #                                   kernel_size=(1, 1),\n",
    "  #                                   padding='same',\n",
    "  #                                   activation=None,\n",
    "  #                                   kernel_initializer=self.initializer,\n",
    "  #                                   kernel_regularizer=self.regularizer,\n",
    "  #                                   name=\"output_conv\")(net)\n",
    "\n",
    "\n",
    "  #   output = net\n",
    "\n",
    "\n",
    "\n",
    "  #       # 保存编码器部分的最终表示，可以用于其他任务或分析\n",
    "  #   with tf.name_scope(\"representation\"):\n",
    "  #       self.representation = convs[-1]\n",
    "\n",
    "  #   # logits 是最终的网络输出，未经过激活函数\n",
    "  #   with tf.name_scope(\"logits\"):\n",
    "  #       self.logits = output  # 假设 'net' 是模型最终输出的张量（output）\n",
    "\n",
    "  #       # 记录 logits 的直方图摘要，以便在 TensorBoard 中可视化\n",
    "  #       tf.summary.histogram(\"logits\", self.logits, step=self.global_step)\n",
    "\n",
    "  #   # 计算最终预测结果，使用 softmax 将 logits 转换为概率分布\n",
    "  #   with tf.name_scope(\"preds\"):\n",
    "  #       self.preds = tf.nn.softmax(self.logits)\n",
    "\n",
    "  #       # 记录预测结果的直方图摘要，以便在 TensorBoard 中可视化\n",
    "  #       tf.summary.histogram(\"preds\", self.preds, step=self.global_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"UNet_3Plus\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " input_layer (InputLayer)    [(None, 320, 320, 1)]        0         []                            \n",
      "                                                                                                  \n",
      " conv2d_2 (Conv2D)           (None, 320, 320, 64)         640       ['input_layer[0][0]']         \n",
      "                                                                                                  \n",
      " batch_normalization_2 (Bat  (None, 320, 320, 64)         256       ['conv2d_2[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_2 (TFOpLambda)   (None, 320, 320, 64)         0         ['batch_normalization_2[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_3 (Conv2D)           (None, 320, 320, 64)         36928     ['tf.nn.relu_2[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_3 (Bat  (None, 320, 320, 64)         256       ['conv2d_3[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_3 (TFOpLambda)   (None, 320, 320, 64)         0         ['batch_normalization_3[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPoolin  (None, 160, 160, 64)         0         ['tf.nn.relu_3[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_4 (Conv2D)           (None, 160, 160, 128)        73856     ['max_pooling2d_1[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_4 (Bat  (None, 160, 160, 128)        512       ['conv2d_4[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_4 (TFOpLambda)   (None, 160, 160, 128)        0         ['batch_normalization_4[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_5 (Conv2D)           (None, 160, 160, 128)        147584    ['tf.nn.relu_4[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_5 (Bat  (None, 160, 160, 128)        512       ['conv2d_5[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_5 (TFOpLambda)   (None, 160, 160, 128)        0         ['batch_normalization_5[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPoolin  (None, 80, 80, 128)          0         ['tf.nn.relu_5[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_6 (Conv2D)           (None, 80, 80, 256)          295168    ['max_pooling2d_2[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_6 (Bat  (None, 80, 80, 256)          1024      ['conv2d_6[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_6 (TFOpLambda)   (None, 80, 80, 256)          0         ['batch_normalization_6[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_7 (Conv2D)           (None, 80, 80, 256)          590080    ['tf.nn.relu_6[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_7 (Bat  (None, 80, 80, 256)          1024      ['conv2d_7[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_7 (TFOpLambda)   (None, 80, 80, 256)          0         ['batch_normalization_7[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_3 (MaxPoolin  (None, 40, 40, 256)          0         ['tf.nn.relu_7[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_8 (Conv2D)           (None, 40, 40, 512)          1180160   ['max_pooling2d_3[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_8 (Bat  (None, 40, 40, 512)          2048      ['conv2d_8[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_8 (TFOpLambda)   (None, 40, 40, 512)          0         ['batch_normalization_8[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_9 (Conv2D)           (None, 40, 40, 512)          2359808   ['tf.nn.relu_8[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_9 (Bat  (None, 40, 40, 512)          2048      ['conv2d_9[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_9 (TFOpLambda)   (None, 40, 40, 512)          0         ['batch_normalization_9[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_4 (MaxPoolin  (None, 20, 20, 512)          0         ['tf.nn.relu_9[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_10 (Conv2D)          (None, 20, 20, 1024)         4719616   ['max_pooling2d_4[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_10 (Ba  (None, 20, 20, 1024)         4096      ['conv2d_10[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_10 (TFOpLambda)  (None, 20, 20, 1024)         0         ['batch_normalization_10[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " conv2d_11 (Conv2D)          (None, 20, 20, 1024)         9438208   ['tf.nn.relu_10[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_11 (Ba  (None, 20, 20, 1024)         4096      ['conv2d_11[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_11 (TFOpLambda)  (None, 20, 20, 1024)         0         ['batch_normalization_11[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " max_pooling2d_5 (MaxPoolin  (None, 40, 40, 64)           0         ['tf.nn.relu_3[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " max_pooling2d_6 (MaxPoolin  (None, 40, 40, 128)          0         ['tf.nn.relu_5[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " max_pooling2d_7 (MaxPoolin  (None, 40, 40, 256)          0         ['tf.nn.relu_7[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d (UpSampling2  (None, 40, 40, 1024)         0         ['tf.nn.relu_11[0][0]']       \n",
      " D)                                                                                               \n",
      "                                                                                                  \n",
      " conv2d_12 (Conv2D)          (None, 40, 40, 64)           36928     ['max_pooling2d_5[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_13 (Conv2D)          (None, 40, 40, 64)           73792     ['max_pooling2d_6[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_14 (Conv2D)          (None, 40, 40, 64)           147520    ['max_pooling2d_7[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_15 (Conv2D)          (None, 40, 40, 64)           294976    ['tf.nn.relu_9[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_16 (Conv2D)          (None, 40, 40, 64)           589888    ['up_sampling2d[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_12 (Ba  (None, 40, 40, 64)           256       ['conv2d_12[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_13 (Ba  (None, 40, 40, 64)           256       ['conv2d_13[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_14 (Ba  (None, 40, 40, 64)           256       ['conv2d_14[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_15 (Ba  (None, 40, 40, 64)           256       ['conv2d_15[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_16 (Ba  (None, 40, 40, 64)           256       ['conv2d_16[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_12 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_12[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_13 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_13[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_14 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_14[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_15 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_15[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_16 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_16[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)   (None, 40, 40, 320)          0         ['tf.nn.relu_12[0][0]',       \n",
      "                                                                     'tf.nn.relu_13[0][0]',       \n",
      "                                                                     'tf.nn.relu_14[0][0]',       \n",
      "                                                                     'tf.nn.relu_15[0][0]',       \n",
      "                                                                     'tf.nn.relu_16[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_17 (Conv2D)          (None, 40, 40, 320)          921920    ['concatenate[0][0]']         \n",
      "                                                                                                  \n",
      " batch_normalization_17 (Ba  (None, 40, 40, 320)          1280      ['conv2d_17[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_17 (TFOpLambda)  (None, 40, 40, 320)          0         ['batch_normalization_17[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " max_pooling2d_8 (MaxPoolin  (None, 80, 80, 64)           0         ['tf.nn.relu_3[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " max_pooling2d_9 (MaxPoolin  (None, 80, 80, 128)          0         ['tf.nn.relu_5[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_1 (UpSamplin  (None, 80, 80, 320)          0         ['tf.nn.relu_17[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_2 (UpSamplin  (None, 80, 80, 1024)         0         ['tf.nn.relu_11[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_18 (Conv2D)          (None, 80, 80, 64)           36928     ['max_pooling2d_8[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_19 (Conv2D)          (None, 80, 80, 64)           73792     ['max_pooling2d_9[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_20 (Conv2D)          (None, 80, 80, 64)           147520    ['tf.nn.relu_7[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_21 (Conv2D)          (None, 80, 80, 64)           184384    ['up_sampling2d_1[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_22 (Conv2D)          (None, 80, 80, 64)           589888    ['up_sampling2d_2[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_18 (Ba  (None, 80, 80, 64)           256       ['conv2d_18[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_19 (Ba  (None, 80, 80, 64)           256       ['conv2d_19[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_20 (Ba  (None, 80, 80, 64)           256       ['conv2d_20[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_21 (Ba  (None, 80, 80, 64)           256       ['conv2d_21[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_22 (Ba  (None, 80, 80, 64)           256       ['conv2d_22[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_18 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_18[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_19 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_19[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_20 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_20[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_21 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_21[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_22 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_22[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate  (None, 80, 80, 320)          0         ['tf.nn.relu_18[0][0]',       \n",
      " )                                                                   'tf.nn.relu_19[0][0]',       \n",
      "                                                                     'tf.nn.relu_20[0][0]',       \n",
      "                                                                     'tf.nn.relu_21[0][0]',       \n",
      "                                                                     'tf.nn.relu_22[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_23 (Conv2D)          (None, 80, 80, 320)          921920    ['concatenate_1[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_23 (Ba  (None, 80, 80, 320)          1280      ['conv2d_23[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_23 (TFOpLambda)  (None, 80, 80, 320)          0         ['batch_normalization_23[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " max_pooling2d_10 (MaxPooli  (None, 160, 160, 64)         0         ['tf.nn.relu_3[0][0]']        \n",
      " ng2D)                                                                                            \n",
      "                                                                                                  \n",
      " up_sampling2d_3 (UpSamplin  (None, 160, 160, 320)        0         ['tf.nn.relu_23[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_4 (UpSamplin  (None, 160, 160, 320)        0         ['tf.nn.relu_17[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_5 (UpSamplin  (None, 160, 160, 1024)       0         ['tf.nn.relu_11[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_24 (Conv2D)          (None, 160, 160, 64)         36928     ['max_pooling2d_10[0][0]']    \n",
      "                                                                                                  \n",
      " conv2d_25 (Conv2D)          (None, 160, 160, 64)         73792     ['tf.nn.relu_5[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_26 (Conv2D)          (None, 160, 160, 64)         184384    ['up_sampling2d_3[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_27 (Conv2D)          (None, 160, 160, 64)         184384    ['up_sampling2d_4[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_28 (Conv2D)          (None, 160, 160, 64)         589888    ['up_sampling2d_5[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_24 (Ba  (None, 160, 160, 64)         256       ['conv2d_24[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_25 (Ba  (None, 160, 160, 64)         256       ['conv2d_25[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_26 (Ba  (None, 160, 160, 64)         256       ['conv2d_26[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_27 (Ba  (None, 160, 160, 64)         256       ['conv2d_27[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_28 (Ba  (None, 160, 160, 64)         256       ['conv2d_28[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_24 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_24[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_25 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_25[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_26 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_26[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_27 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_27[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_28 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_28[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate  (None, 160, 160, 320)        0         ['tf.nn.relu_24[0][0]',       \n",
      " )                                                                   'tf.nn.relu_25[0][0]',       \n",
      "                                                                     'tf.nn.relu_26[0][0]',       \n",
      "                                                                     'tf.nn.relu_27[0][0]',       \n",
      "                                                                     'tf.nn.relu_28[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_29 (Conv2D)          (None, 160, 160, 320)        921920    ['concatenate_2[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_29 (Ba  (None, 160, 160, 320)        1280      ['conv2d_29[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_29 (TFOpLambda)  (None, 160, 160, 320)        0         ['batch_normalization_29[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " up_sampling2d_6 (UpSamplin  (None, 320, 320, 320)        0         ['tf.nn.relu_29[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_7 (UpSamplin  (None, 320, 320, 320)        0         ['tf.nn.relu_23[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_8 (UpSamplin  (None, 320, 320, 320)        0         ['tf.nn.relu_17[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_9 (UpSamplin  (None, 320, 320, 1024)       0         ['tf.nn.relu_11[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_30 (Conv2D)          (None, 320, 320, 64)         36928     ['tf.nn.relu_3[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_31 (Conv2D)          (None, 320, 320, 64)         184384    ['up_sampling2d_6[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_32 (Conv2D)          (None, 320, 320, 64)         184384    ['up_sampling2d_7[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_33 (Conv2D)          (None, 320, 320, 64)         184384    ['up_sampling2d_8[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_34 (Conv2D)          (None, 320, 320, 64)         589888    ['up_sampling2d_9[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_30 (Ba  (None, 320, 320, 64)         256       ['conv2d_30[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_31 (Ba  (None, 320, 320, 64)         256       ['conv2d_31[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_32 (Ba  (None, 320, 320, 64)         256       ['conv2d_32[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_33 (Ba  (None, 320, 320, 64)         256       ['conv2d_33[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_34 (Ba  (None, 320, 320, 64)         256       ['conv2d_34[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_30 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_30[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_31 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_31[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_32 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_32[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_33 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_33[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_34 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_34[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate  (None, 320, 320, 320)        0         ['tf.nn.relu_30[0][0]',       \n",
      " )                                                                   'tf.nn.relu_31[0][0]',       \n",
      "                                                                     'tf.nn.relu_32[0][0]',       \n",
      "                                                                     'tf.nn.relu_33[0][0]',       \n",
      "                                                                     'tf.nn.relu_34[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_35 (Conv2D)          (None, 320, 320, 320)        921920    ['concatenate_3[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_35 (Ba  (None, 320, 320, 320)        1280      ['conv2d_35[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_35 (TFOpLambda)  (None, 320, 320, 320)        0         ['batch_normalization_35[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " conv2d_36 (Conv2D)          (None, 320, 320, 1)          2881      ['tf.nn.relu_35[0][0]']       \n",
      "                                                                                                  \n",
      " tf.nn.softmax (TFOpLambda)  (None, 320, 320, 1)          0         ['conv2d_36[0][0]']           \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 26983681 (102.93 MB)\n",
      "Trainable params: 26970625 (102.88 MB)\n",
      "Non-trainable params: 13056 (51.00 KB)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "UNet3+ base model\n",
    "\"\"\"\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as k\n",
    "\n",
    "def conv_block(x, kernels, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "               is_bn=True, is_relu=True, n=2):\n",
    "    \"\"\" Custom function for conv2d:\n",
    "        Apply  3*3 convolutions with BN and relu.\n",
    "    \"\"\"\n",
    "    for i in range(1, n + 1):\n",
    "        x = k.layers.Conv2D(filters=kernels, kernel_size=kernel_size,\n",
    "                            padding=padding, strides=strides,\n",
    "                            kernel_regularizer=tf.keras.regularizers.l2(1e-4),\n",
    "                            kernel_initializer=k.initializers.he_normal(seed=5))(x)\n",
    "        if is_bn:\n",
    "            x = k.layers.BatchNormalization()(x)\n",
    "        if is_relu:\n",
    "            x = k.activations.relu(x)\n",
    "\n",
    "    return x\n",
    "\n",
    "\n",
    "\n",
    "def unet3plus(input_shape, output_channels):\n",
    "    \"\"\" UNet3+ base model \"\"\"\n",
    "    filters = [64, 128, 256, 512, 1024]\n",
    "\n",
    "    input_layer = k.layers.Input(\n",
    "        shape=input_shape,\n",
    "        name=\"input_layer\"\n",
    "    )  # 320*320*3\n",
    "\n",
    "    \"\"\" Encoder\"\"\"\n",
    "    # block 1\n",
    "    e1 = conv_block(input_layer, filters[0])  # 320*320*64\n",
    "\n",
    "    # block 2\n",
    "    e2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 160*160*64\n",
    "    e2 = conv_block(e2, filters[1])  # 160*160*128\n",
    "\n",
    "    # block 3\n",
    "    e3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 80*80*128\n",
    "    e3 = conv_block(e3, filters[2])  # 80*80*256\n",
    "\n",
    "    # block 4\n",
    "    e4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 40*40*256\n",
    "    e4 = conv_block(e4, filters[3])  # 40*40*512\n",
    "\n",
    "    # block 5\n",
    "    # bottleneck layer\n",
    "    e5 = k.layers.MaxPool2D(pool_size=(2, 2))(e4)  # 20*20*512\n",
    "    e5 = conv_block(e5, filters[4])  # 20*20*1024   \n",
    "\n",
    "    \"\"\" Decoder \"\"\"\n",
    "    cat_channels = filters[0]\n",
    "    cat_blocks = len(filters)\n",
    "    upsample_channels = cat_blocks * cat_channels\n",
    "\n",
    "    \"\"\" d4 \"\"\"\n",
    "    e1_d4 = k.layers.MaxPool2D(pool_size=(8, 8))(e1)  # 320*320*64  --> 40*40*64\n",
    "    e1_d4 = conv_block(e1_d4, cat_channels, n=1)  # 320*320*64  --> 40*40*64\n",
    "\n",
    "    e2_d4 = k.layers.MaxPool2D(pool_size=(4, 4))(e2)  # 160*160*128 --> 40*40*128\n",
    "    e2_d4 = conv_block(e2_d4, cat_channels, n=1)  # 160*160*128 --> 40*40*64\n",
    "\n",
    "    e3_d4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 80*80*256  --> 40*40*256\n",
    "    e3_d4 = conv_block(e3_d4, cat_channels, n=1)  # 80*80*256  --> 40*40*64\n",
    "\n",
    "    e4_d4 = conv_block(e4, cat_channels, n=1)  # 40*40*512  --> 40*40*64\n",
    "\n",
    "    e5_d4 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(e5)  # 80*80*256  --> 40*40*256\n",
    "    e5_d4 = conv_block(e5_d4, cat_channels, n=1)  # 20*20*1024  --> 20*20*64\n",
    "\n",
    "    d4 = k.layers.concatenate([e1_d4, e2_d4, e3_d4, e4_d4, e5_d4])\n",
    "    d4 = conv_block(d4, upsample_channels, n=1)  # 40*40*320  --> 40*40*320\n",
    "\n",
    "    \"\"\" d3 \"\"\"\n",
    "    e1_d3 = k.layers.MaxPool2D(pool_size=(4, 4))(e1)  # 320*320*64 --> 80*80*64\n",
    "    e1_d3 = conv_block(e1_d3, cat_channels, n=1)  # 80*80*64 --> 80*80*64\n",
    "\n",
    "    e2_d3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 160*160*256 --> 80*80*256\n",
    "    e2_d3 = conv_block(e2_d3, cat_channels, n=1)  # 80*80*256 --> 80*80*64\n",
    "\n",
    "    e3_d3 = conv_block(e3, cat_channels, n=1)  # 80*80*512 --> 80*80*64\n",
    "\n",
    "    e4_d3 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d4)  # 40*40*320 --> 80*80*320\n",
    "    e4_d3 = conv_block(e4_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64\n",
    "\n",
    "    e5_d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(e5)  # 20*20*320 --> 80*80*320\n",
    "    e5_d3 = conv_block(e5_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64\n",
    "\n",
    "    d3 = k.layers.concatenate([e1_d3, e2_d3, e3_d3, e4_d3, e5_d3])\n",
    "    d3 = conv_block(d3, upsample_channels, n=1)  # 80*80*320 --> 80*80*320\n",
    "\n",
    "    \"\"\" d2 \"\"\"\n",
    "    e1_d2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 320*320*64 --> 160*160*64\n",
    "    e1_d2 = conv_block(e1_d2, cat_channels, n=1)  # 160*160*64 --> 160*160*64\n",
    "\n",
    "    e2_d2 = conv_block(e2, cat_channels, n=1)  # 160*160*256 --> 160*160*64\n",
    "\n",
    "    d3_d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d3)  # 80*80*320 --> 160*160*320\n",
    "    d3_d2 = conv_block(d3_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    d4_d2 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d4)  # 40*40*320 --> 160*160*320\n",
    "    d4_d2 = conv_block(d4_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    e5_d2 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(e5)  # 20*20*320 --> 160*160*320\n",
    "    e5_d2 = conv_block(e5_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    d2 = k.layers.concatenate([e1_d2, e2_d2, d3_d2, d4_d2, e5_d2])\n",
    "    d2 = conv_block(d2, upsample_channels, n=1)  # 160*160*320 --> 160*160*320\n",
    "\n",
    "    \"\"\" d1 \"\"\"\n",
    "    e1_d1 = conv_block(e1, cat_channels, n=1)  # 320*320*64 --> 320*320*64\n",
    "\n",
    "    d2_d1 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)  # 160*160*320 --> 320*320*320\n",
    "    d2_d1 = conv_block(d2_d1, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    d3_d1 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)  # 80*80*320 --> 320*320*320\n",
    "    d3_d1 = conv_block(d3_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64\n",
    "\n",
    "    d4_d1 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)  # 40*40*320 --> 320*320*320\n",
    "    d4_d1 = conv_block(d4_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64\n",
    "\n",
    "    e5_d1 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)  # 20*20*320 --> 320*320*320\n",
    "    e5_d1 = conv_block(e5_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64\n",
    "\n",
    "    d1 = k.layers.concatenate([e1_d1, d2_d1, d3_d1, d4_d1, e5_d1, ])\n",
    "    d1 = conv_block(d1, upsample_channels, n=1)  # 320*320*320 --> 320*320*320\n",
    "\n",
    "    # last layer does not have batchnorm and relu\n",
    "    d = conv_block(d1, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "\n",
    "    output = k.activations.softmax(d)\n",
    "\n",
    "    return tf.keras.Model(inputs=input_layer, outputs=[output], name='UNet_3Plus')\n",
    "\n",
    "\n",
    "def tiny_unet3plus(input_shape, output_channels, training):\n",
    "    \"\"\" Sample model only for testing during development \"\"\"\n",
    "    filters = [64, 128, 256, 512, 1024]\n",
    "\n",
    "    input_layer = k.layers.Input(shape=input_shape, name=\"input_layer\")  # 320*320*3\n",
    "\n",
    "    \"\"\" Encoder\"\"\"\n",
    "    # block 1\n",
    "    e1 = conv_block(input_layer, filters[0] // 2)  # 320*320*64\n",
    "    e1 = conv_block(e1, filters[0] // 2)  # 320*320*64\n",
    "\n",
    "    # last layer does not have batch norm and relu\n",
    "    d = conv_block(e1, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "    output = k.activations.softmax(d, )\n",
    "\n",
    "    if training:\n",
    "        e2 = conv_block(e1, filters[0] // 2)  # 320*320*64\n",
    "        d2 = conv_block(e2, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "        output2 = k.activations.softmax(d2)\n",
    "        #output2=k.activations.sigmoid(d2)\n",
    "        #output2=k.activations.tanh(d2)\n",
    "        return tf.keras.Model(inputs=input_layer, outputs=[output, output2], name='UNet3Plus')\n",
    "    else:\n",
    "        return tf.keras.Model(inputs=input_layer, outputs=[output], name='UNet3Plus')\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    \"\"\"## Model Compilation\"\"\"\n",
    "    INPUT_SHAPE = [320, 320, 1]\n",
    "    OUTPUT_CHANNELS = 1\n",
    "\n",
    "    unet_3P = unet3plus(INPUT_SHAPE, OUTPUT_CHANNELS)\n",
    "    unet_3P.summary()\n",
    "\n",
    "    # tf.keras.utils.plot_model(unet_3P, show_layer_names=True, show_shapes=True)\n",
    "\n",
    "    # unet_3P.save(\"unet_3P.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"UNet3Plus_DeepSup\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " input_layer (InputLayer)    [(None, 320, 320, 1)]        0         []                            \n",
      "                                                                                                  \n",
      " conv2d (Conv2D)             (None, 320, 320, 64)         640       ['input_layer[0][0]']         \n",
      "                                                                                                  \n",
      " batch_normalization (Batch  (None, 320, 320, 64)         256       ['conv2d[0][0]']              \n",
      " Normalization)                                                                                   \n",
      "                                                                                                  \n",
      " tf.nn.relu (TFOpLambda)     (None, 320, 320, 64)         0         ['batch_normalization[0][0]'] \n",
      "                                                                                                  \n",
      " conv2d_1 (Conv2D)           (None, 320, 320, 64)         36928     ['tf.nn.relu[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_1 (Bat  (None, 320, 320, 64)         256       ['conv2d_1[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_1 (TFOpLambda)   (None, 320, 320, 64)         0         ['batch_normalization_1[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d (MaxPooling2  (None, 160, 160, 64)         0         ['tf.nn.relu_1[0][0]']        \n",
      " D)                                                                                               \n",
      "                                                                                                  \n",
      " conv2d_2 (Conv2D)           (None, 160, 160, 128)        73856     ['max_pooling2d[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_2 (Bat  (None, 160, 160, 128)        512       ['conv2d_2[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_2 (TFOpLambda)   (None, 160, 160, 128)        0         ['batch_normalization_2[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_3 (Conv2D)           (None, 160, 160, 128)        147584    ['tf.nn.relu_2[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_3 (Bat  (None, 160, 160, 128)        512       ['conv2d_3[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_3 (TFOpLambda)   (None, 160, 160, 128)        0         ['batch_normalization_3[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPoolin  (None, 80, 80, 128)          0         ['tf.nn.relu_3[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_4 (Conv2D)           (None, 80, 80, 256)          295168    ['max_pooling2d_1[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_4 (Bat  (None, 80, 80, 256)          1024      ['conv2d_4[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_4 (TFOpLambda)   (None, 80, 80, 256)          0         ['batch_normalization_4[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_5 (Conv2D)           (None, 80, 80, 256)          590080    ['tf.nn.relu_4[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_5 (Bat  (None, 80, 80, 256)          1024      ['conv2d_5[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_5 (TFOpLambda)   (None, 80, 80, 256)          0         ['batch_normalization_5[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPoolin  (None, 40, 40, 256)          0         ['tf.nn.relu_5[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_6 (Conv2D)           (None, 40, 40, 512)          1180160   ['max_pooling2d_2[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_6 (Bat  (None, 40, 40, 512)          2048      ['conv2d_6[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_6 (TFOpLambda)   (None, 40, 40, 512)          0         ['batch_normalization_6[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_7 (Conv2D)           (None, 40, 40, 512)          2359808   ['tf.nn.relu_6[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_7 (Bat  (None, 40, 40, 512)          2048      ['conv2d_7[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_7 (TFOpLambda)   (None, 40, 40, 512)          0         ['batch_normalization_7[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_3 (MaxPoolin  (None, 20, 20, 512)          0         ['tf.nn.relu_7[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_8 (Conv2D)           (None, 20, 20, 1024)         4719616   ['max_pooling2d_3[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_8 (Bat  (None, 20, 20, 1024)         4096      ['conv2d_8[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_8 (TFOpLambda)   (None, 20, 20, 1024)         0         ['batch_normalization_8[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " conv2d_9 (Conv2D)           (None, 20, 20, 1024)         9438208   ['tf.nn.relu_8[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_9 (Bat  (None, 20, 20, 1024)         4096      ['conv2d_9[0][0]']            \n",
      " chNormalization)                                                                                 \n",
      "                                                                                                  \n",
      " tf.nn.relu_9 (TFOpLambda)   (None, 20, 20, 1024)         0         ['batch_normalization_9[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " max_pooling2d_4 (MaxPoolin  (None, 40, 40, 64)           0         ['tf.nn.relu_1[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " max_pooling2d_5 (MaxPoolin  (None, 40, 40, 128)          0         ['tf.nn.relu_3[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " max_pooling2d_6 (MaxPoolin  (None, 40, 40, 256)          0         ['tf.nn.relu_5[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d (UpSampling2  (None, 40, 40, 1024)         0         ['tf.nn.relu_9[0][0]']        \n",
      " D)                                                                                               \n",
      "                                                                                                  \n",
      " conv2d_10 (Conv2D)          (None, 40, 40, 64)           36928     ['max_pooling2d_4[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_11 (Conv2D)          (None, 40, 40, 64)           73792     ['max_pooling2d_5[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_12 (Conv2D)          (None, 40, 40, 64)           147520    ['max_pooling2d_6[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_13 (Conv2D)          (None, 40, 40, 64)           294976    ['tf.nn.relu_7[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_14 (Conv2D)          (None, 40, 40, 64)           589888    ['up_sampling2d[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_10 (Ba  (None, 40, 40, 64)           256       ['conv2d_10[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_11 (Ba  (None, 40, 40, 64)           256       ['conv2d_11[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_12 (Ba  (None, 40, 40, 64)           256       ['conv2d_12[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_13 (Ba  (None, 40, 40, 64)           256       ['conv2d_13[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_14 (Ba  (None, 40, 40, 64)           256       ['conv2d_14[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_10 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_10[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_11 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_11[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_12 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_12[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_13 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_13[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_14 (TFOpLambda)  (None, 40, 40, 64)           0         ['batch_normalization_14[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)   (None, 40, 40, 320)          0         ['tf.nn.relu_10[0][0]',       \n",
      "                                                                     'tf.nn.relu_11[0][0]',       \n",
      "                                                                     'tf.nn.relu_12[0][0]',       \n",
      "                                                                     'tf.nn.relu_13[0][0]',       \n",
      "                                                                     'tf.nn.relu_14[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_15 (Conv2D)          (None, 40, 40, 320)          921920    ['concatenate[0][0]']         \n",
      "                                                                                                  \n",
      " batch_normalization_15 (Ba  (None, 40, 40, 320)          1280      ['conv2d_15[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_15 (TFOpLambda)  (None, 40, 40, 320)          0         ['batch_normalization_15[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " max_pooling2d_7 (MaxPoolin  (None, 80, 80, 64)           0         ['tf.nn.relu_1[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " max_pooling2d_8 (MaxPoolin  (None, 80, 80, 128)          0         ['tf.nn.relu_3[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_1 (UpSamplin  (None, 80, 80, 320)          0         ['tf.nn.relu_15[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_2 (UpSamplin  (None, 80, 80, 1024)         0         ['tf.nn.relu_9[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_16 (Conv2D)          (None, 80, 80, 64)           36928     ['max_pooling2d_7[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_17 (Conv2D)          (None, 80, 80, 64)           73792     ['max_pooling2d_8[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_18 (Conv2D)          (None, 80, 80, 64)           147520    ['tf.nn.relu_5[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_19 (Conv2D)          (None, 80, 80, 64)           184384    ['up_sampling2d_1[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_20 (Conv2D)          (None, 80, 80, 64)           589888    ['up_sampling2d_2[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_16 (Ba  (None, 80, 80, 64)           256       ['conv2d_16[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_17 (Ba  (None, 80, 80, 64)           256       ['conv2d_17[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_18 (Ba  (None, 80, 80, 64)           256       ['conv2d_18[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_19 (Ba  (None, 80, 80, 64)           256       ['conv2d_19[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_20 (Ba  (None, 80, 80, 64)           256       ['conv2d_20[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_16 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_16[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_17 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_17[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_18 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_18[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_19 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_19[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_20 (TFOpLambda)  (None, 80, 80, 64)           0         ['batch_normalization_20[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate_1 (Concatenate  (None, 80, 80, 320)          0         ['tf.nn.relu_16[0][0]',       \n",
      " )                                                                   'tf.nn.relu_17[0][0]',       \n",
      "                                                                     'tf.nn.relu_18[0][0]',       \n",
      "                                                                     'tf.nn.relu_19[0][0]',       \n",
      "                                                                     'tf.nn.relu_20[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_21 (Conv2D)          (None, 80, 80, 320)          921920    ['concatenate_1[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_21 (Ba  (None, 80, 80, 320)          1280      ['conv2d_21[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_21 (TFOpLambda)  (None, 80, 80, 320)          0         ['batch_normalization_21[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " max_pooling2d_9 (MaxPoolin  (None, 160, 160, 64)         0         ['tf.nn.relu_1[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_3 (UpSamplin  (None, 160, 160, 320)        0         ['tf.nn.relu_21[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_4 (UpSamplin  (None, 160, 160, 320)        0         ['tf.nn.relu_15[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_5 (UpSamplin  (None, 160, 160, 1024)       0         ['tf.nn.relu_9[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_22 (Conv2D)          (None, 160, 160, 64)         36928     ['max_pooling2d_9[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_23 (Conv2D)          (None, 160, 160, 64)         73792     ['tf.nn.relu_3[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_24 (Conv2D)          (None, 160, 160, 64)         184384    ['up_sampling2d_3[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_25 (Conv2D)          (None, 160, 160, 64)         184384    ['up_sampling2d_4[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_26 (Conv2D)          (None, 160, 160, 64)         589888    ['up_sampling2d_5[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_22 (Ba  (None, 160, 160, 64)         256       ['conv2d_22[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_23 (Ba  (None, 160, 160, 64)         256       ['conv2d_23[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_24 (Ba  (None, 160, 160, 64)         256       ['conv2d_24[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_25 (Ba  (None, 160, 160, 64)         256       ['conv2d_25[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_26 (Ba  (None, 160, 160, 64)         256       ['conv2d_26[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_22 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_22[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_23 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_23[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_24 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_24[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_25 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_25[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_26 (TFOpLambda)  (None, 160, 160, 64)         0         ['batch_normalization_26[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate_2 (Concatenate  (None, 160, 160, 320)        0         ['tf.nn.relu_22[0][0]',       \n",
      " )                                                                   'tf.nn.relu_23[0][0]',       \n",
      "                                                                     'tf.nn.relu_24[0][0]',       \n",
      "                                                                     'tf.nn.relu_25[0][0]',       \n",
      "                                                                     'tf.nn.relu_26[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_27 (Conv2D)          (None, 160, 160, 320)        921920    ['concatenate_2[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_27 (Ba  (None, 160, 160, 320)        1280      ['conv2d_27[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_27 (TFOpLambda)  (None, 160, 160, 320)        0         ['batch_normalization_27[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " up_sampling2d_6 (UpSamplin  (None, 320, 320, 320)        0         ['tf.nn.relu_27[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_7 (UpSamplin  (None, 320, 320, 320)        0         ['tf.nn.relu_21[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_8 (UpSamplin  (None, 320, 320, 320)        0         ['tf.nn.relu_15[0][0]']       \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " up_sampling2d_9 (UpSamplin  (None, 320, 320, 1024)       0         ['tf.nn.relu_9[0][0]']        \n",
      " g2D)                                                                                             \n",
      "                                                                                                  \n",
      " conv2d_28 (Conv2D)          (None, 320, 320, 64)         36928     ['tf.nn.relu_1[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_29 (Conv2D)          (None, 320, 320, 64)         184384    ['up_sampling2d_6[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_30 (Conv2D)          (None, 320, 320, 64)         184384    ['up_sampling2d_7[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_31 (Conv2D)          (None, 320, 320, 64)         184384    ['up_sampling2d_8[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_32 (Conv2D)          (None, 320, 320, 64)         589888    ['up_sampling2d_9[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_28 (Ba  (None, 320, 320, 64)         256       ['conv2d_28[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_29 (Ba  (None, 320, 320, 64)         256       ['conv2d_29[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_30 (Ba  (None, 320, 320, 64)         256       ['conv2d_30[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_31 (Ba  (None, 320, 320, 64)         256       ['conv2d_31[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " batch_normalization_32 (Ba  (None, 320, 320, 64)         256       ['conv2d_32[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_28 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_28[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_29 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_29[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_30 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_30[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_31 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_31[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.nn.relu_32 (TFOpLambda)  (None, 320, 320, 64)         0         ['batch_normalization_32[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " concatenate_3 (Concatenate  (None, 320, 320, 320)        0         ['tf.nn.relu_28[0][0]',       \n",
      " )                                                                   'tf.nn.relu_29[0][0]',       \n",
      "                                                                     'tf.nn.relu_30[0][0]',       \n",
      "                                                                     'tf.nn.relu_31[0][0]',       \n",
      "                                                                     'tf.nn.relu_32[0][0]']       \n",
      "                                                                                                  \n",
      " conv2d_33 (Conv2D)          (None, 320, 320, 320)        921920    ['concatenate_3[0][0]']       \n",
      "                                                                                                  \n",
      " batch_normalization_33 (Ba  (None, 320, 320, 320)        1280      ['conv2d_33[0][0]']           \n",
      " tchNormalization)                                                                                \n",
      "                                                                                                  \n",
      " tf.nn.relu_33 (TFOpLambda)  (None, 320, 320, 320)        0         ['batch_normalization_33[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " conv2d_34 (Conv2D)          (None, 320, 320, 1)          2881      ['tf.nn.relu_33[0][0]']       \n",
      "                                                                                                  \n",
      " tf.nn.softmax (TFOpLambda)  (None, 320, 320, 1)          0         ['conv2d_34[0][0]']           \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 26983681 (102.93 MB)\n",
      "Trainable params: 26970625 (102.88 MB)\n",
      "Non-trainable params: 13056 (51.00 KB)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "UNet3+ with Deep Supervision\n",
    "\"\"\"\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as k\n",
    "\n",
    "def conv_block(x, kernels, kernel_size=(3, 3), strides=(1, 1), padding='same',\n",
    "               is_bn=True, is_relu=True, n=2):\n",
    "    \"\"\" Custom function for conv2d:\n",
    "        Apply  3*3 convolutions with BN and relu.\n",
    "    \"\"\"\n",
    "    for i in range(1, n + 1):\n",
    "        x = k.layers.Conv2D(filters=kernels, kernel_size=kernel_size,\n",
    "                            padding=padding, strides=strides,\n",
    "                            kernel_regularizer=tf.keras.regularizers.l2(1e-4),\n",
    "                            kernel_initializer=k.initializers.he_normal(seed=5))(x)\n",
    "        if is_bn:\n",
    "            x = k.layers.BatchNormalization()(x)\n",
    "        if is_relu:\n",
    "            x = k.activations.relu(x)\n",
    "\n",
    "    return x\n",
    "\n",
    "\n",
    "def unet3plus_deepsup(input_shape, output_channels, training=False):\n",
    "    \"\"\" UNet_3Plus with Deep Supervision \"\"\"\n",
    "    filters = [64, 128, 256, 512, 1024]\n",
    "\n",
    "    input_layer = k.layers.Input(shape=input_shape, name=\"input_layer\")  # 320*320*3\n",
    "\n",
    "    \"\"\" Encoder\"\"\"\n",
    "    # block 1\n",
    "    e1 = conv_block(input_layer, filters[0])  # 320*320*64\n",
    "\n",
    "    # block 2\n",
    "    e2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 160*160*64\n",
    "    e2 = conv_block(e2, filters[1])  # 160*160*128\n",
    "\n",
    "    # block 3\n",
    "    e3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 80*80*128\n",
    "    e3 = conv_block(e3, filters[2])  # 80*80*256\n",
    "\n",
    "    # block 4\n",
    "    e4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 40*40*256\n",
    "    e4 = conv_block(e4, filters[3])  # 40*40*512\n",
    "\n",
    "    # block 5\n",
    "    # bottleneck layer\n",
    "    e5 = k.layers.MaxPool2D(pool_size=(2, 2))(e4)  # 20*20*512\n",
    "    e5 = conv_block(e5, filters[4])  # 20*20*1024\n",
    "\n",
    "    \"\"\" Decoder \"\"\"\n",
    "    cat_channels = filters[0]\n",
    "    cat_blocks = len(filters)\n",
    "    upsample_channels = cat_blocks * cat_channels\n",
    "\n",
    "    \"\"\" d4 \"\"\"\n",
    "    e1_d4 = k.layers.MaxPool2D(pool_size=(8, 8))(e1)  # 320*320*64  --> 40*40*64\n",
    "    e1_d4 = conv_block(e1_d4, cat_channels, n=1)  # 320*320*64  --> 40*40*64\n",
    "\n",
    "    e2_d4 = k.layers.MaxPool2D(pool_size=(4, 4))(e2)  # 160*160*128 --> 40*40*128\n",
    "    e2_d4 = conv_block(e2_d4, cat_channels, n=1)  # 160*160*128 --> 40*40*64\n",
    "\n",
    "    e3_d4 = k.layers.MaxPool2D(pool_size=(2, 2))(e3)  # 80*80*256  --> 40*40*256\n",
    "    e3_d4 = conv_block(e3_d4, cat_channels, n=1)  # 80*80*256  --> 40*40*64\n",
    "\n",
    "    e4_d4 = conv_block(e4, cat_channels, n=1)  # 40*40*512  --> 40*40*64\n",
    "\n",
    "    e5_d4 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(e5)  # 80*80*256  --> 40*40*256\n",
    "    e5_d4 = conv_block(e5_d4, cat_channels, n=1)  # 20*20*1024  --> 20*20*64\n",
    "\n",
    "    d4 = k.layers.concatenate([e1_d4, e2_d4, e3_d4, e4_d4, e5_d4])\n",
    "    d4 = conv_block(d4, upsample_channels, n=1)  # 40*40*320  --> 40*40*320\n",
    "\n",
    "    \"\"\" d3 \"\"\"\n",
    "    e1_d3 = k.layers.MaxPool2D(pool_size=(4, 4))(e1)  # 320*320*64 --> 80*80*64\n",
    "    e1_d3 = conv_block(e1_d3, cat_channels, n=1)  # 80*80*64 --> 80*80*64\n",
    "\n",
    "    e2_d3 = k.layers.MaxPool2D(pool_size=(2, 2))(e2)  # 160*160*256 --> 80*80*256\n",
    "    e2_d3 = conv_block(e2_d3, cat_channels, n=1)  # 80*80*256 --> 80*80*64\n",
    "\n",
    "    e3_d3 = conv_block(e3, cat_channels, n=1)  # 80*80*512 --> 80*80*64\n",
    "\n",
    "    e4_d3 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d4)  # 40*40*320 --> 80*80*320\n",
    "    e4_d3 = conv_block(e4_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64\n",
    "\n",
    "    e5_d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(e5)  # 20*20*320 --> 80*80*320\n",
    "    e5_d3 = conv_block(e5_d3, cat_channels, n=1)  # 80*80*320 --> 80*80*64\n",
    "\n",
    "    d3 = k.layers.concatenate([e1_d3, e2_d3, e3_d3, e4_d3, e5_d3])\n",
    "    d3 = conv_block(d3, upsample_channels, n=1)  # 80*80*320 --> 80*80*320\n",
    "\n",
    "    \"\"\" d2 \"\"\"\n",
    "    e1_d2 = k.layers.MaxPool2D(pool_size=(2, 2))(e1)  # 320*320*64 --> 160*160*64\n",
    "    e1_d2 = conv_block(e1_d2, cat_channels, n=1)  # 160*160*64 --> 160*160*64\n",
    "\n",
    "    e2_d2 = conv_block(e2, cat_channels, n=1)  # 160*160*256 --> 160*160*64\n",
    "\n",
    "    d3_d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d3)  # 80*80*320 --> 160*160*320\n",
    "    d3_d2 = conv_block(d3_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    d4_d2 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d4)  # 40*40*320 --> 160*160*320\n",
    "    d4_d2 = conv_block(d4_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    e5_d2 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(e5)  # 20*20*320 --> 160*160*320\n",
    "    e5_d2 = conv_block(e5_d2, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    d2 = k.layers.concatenate([e1_d2, e2_d2, d3_d2, d4_d2, e5_d2])\n",
    "    d2 = conv_block(d2, upsample_channels, n=1)  # 160*160*320 --> 160*160*320\n",
    "\n",
    "    \"\"\" d1 \"\"\"\n",
    "    e1_d1 = conv_block(e1, cat_channels, n=1)  # 320*320*64 --> 320*320*64\n",
    "\n",
    "    d2_d1 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)  # 160*160*320 --> 320*320*320\n",
    "    d2_d1 = conv_block(d2_d1, cat_channels, n=1)  # 160*160*320 --> 160*160*64\n",
    "\n",
    "    d3_d1 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)  # 80*80*320 --> 320*320*320\n",
    "    d3_d1 = conv_block(d3_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64\n",
    "\n",
    "    d4_d1 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)  # 40*40*320 --> 320*320*320\n",
    "    d4_d1 = conv_block(d4_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64\n",
    "\n",
    "    e5_d1 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)  # 20*20*320 --> 320*320*320\n",
    "    e5_d1 = conv_block(e5_d1, cat_channels, n=1)  # 320*320*320 --> 320*320*64\n",
    "\n",
    "    d1 = k.layers.concatenate([e1_d1, d2_d1, d3_d1, d4_d1, e5_d1, ])\n",
    "    d1 = conv_block(d1, upsample_channels, n=1)  # 320*320*320 --> 320*320*320\n",
    "\n",
    "    # last layer does not have batch norm and relu\n",
    "    d1 = conv_block(d1, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "    d1 = k.activations.softmax(d1)\n",
    "\n",
    "    \"\"\" Deep Supervision Part\"\"\"\n",
    "    if training:\n",
    "        d2 = conv_block(d2, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "        d3 = conv_block(d3, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "        d4 = conv_block(d4, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "        e5 = conv_block(e5, output_channels, n=1, is_bn=False, is_relu=False)\n",
    "\n",
    "        # d1 = no need for up sampling\n",
    "        d2 = k.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(d2)#80*80*1\n",
    "        d3 = k.layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(d3)#20*20*1\n",
    "        d4 = k.layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(d4)#5*5*1\n",
    "        e5 = k.layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(e5)#1*1*1\n",
    "\n",
    "        d2 = k.activations.softmax(d2)\n",
    "        d3 = k.activations.softmax(d3)\n",
    "        d4 = k.activations.softmax(d4)\n",
    "        e5 = k.activations.softmax(e5)\n",
    "\n",
    "    if training:\n",
    "        return tf.keras.Model(inputs=input_layer, outputs=[d1, d2, d3, d4, e5], name='UNet3Plus_DeepSup')\n",
    "    else:\n",
    "        return tf.keras.Model(inputs=input_layer, outputs=[d1, ], name='UNet3Plus_DeepSup')\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    \"\"\"## Model Compilation\"\"\"\n",
    "    INPUT_SHAPE = [320, 320, 1]\n",
    "    OUTPUT_CHANNELS = 1\n",
    "\n",
    "    unet_3P = unet3plus_deepsup(INPUT_SHAPE, OUTPUT_CHANNELS)\n",
    "    unet_3P.summary()\n",
    "\n",
    "    # tf.keras.utils.plot_model(unet_3P, show_layer_names=True, show_shapes=True)\n",
    "\n",
    "    # unet_3P.save(\"unet_3P.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "    #卷积块\n",
    "  def conv_block(self,net,filters_root,n=2):\n",
    "    for i in range(1, n+1):\n",
    "        net = tf.keras.layers.Conv2D(\n",
    "        filters=filters_root, \n",
    "        kernel_size=self.kernel_size,\n",
    "        padding='same', strides=(1,1),\n",
    "        kernel_regularizer=self.regularizer ,\n",
    "        kernel_initializer=self.initializer)(net)\n",
    "\n",
    "        net = tf.keras.layers.BatchNormalization()(net)\n",
    "        net = tf.keras.activations.relu(net)\n",
    "\n",
    "        net= tf.keras.layers.Dropout(rate=self.drop_rate)(net)\n",
    "\n",
    "    return net\n",
    "\n",
    "  def add_prediction_op(self):\n",
    "    logging.info(\"Model: depths {depths}, filters {filters}, \"\n",
    "           \"filter size {kernel_size[0]}x{kernel_size[1]}, \"\n",
    "           \"pool size: {pool_size[0]}x{pool_size[1]}, \"\n",
    "           \"dilation rate: {dilation_rate[0]}x{dilation_rate[1]}\".format(\n",
    "            depths=self.depths,\n",
    "            filters=self.filters_root,\n",
    "            kernel_size=self.kernel_size,\n",
    "            dilation_rate=self.dilation_rate,\n",
    "            pool_size=self.pool_size))\n",
    "\n",
    "      # 如果 weight_decay 大于 0，设置 L2 正则化以防止过拟合，否则不使用正则化\n",
    "    if self.weight_decay > 0:\n",
    "      weight_decay = tf.constant(self.weight_decay, dtype=tf.float32, name=\"weight_constant\")\n",
    "      self.regularizer = tf.keras.regularizers.l2(l=0.5 * (weight_decay))\n",
    "    else:\n",
    "      self.regularizer = None\n",
    "\n",
    "    #权重初始化器，1.0不缩放\n",
    "    self.initializer = tf.compat.v1.keras.initializers.VarianceScaling(scale=1.0, mode=\"fan_avg\", distribution=\"uniform\")\n",
    "\n",
    "\n",
    "    filters_root = [64, 128, 256, 512, 1024]\n",
    "\n",
    "    # net=tf.keras.layers.Input(shape=self.X_shape)# 3000,1,3\n",
    "        # 使用占位符 `self.X` 作为输入张量\n",
    "    with tf.compat.v1.variable_scope(\"Input\"):\n",
    "        net = self.X  # 使用占位符 `self.X` 作为网络的输入\n",
    "\n",
    "    \"\"\" Encoder\"\"\"\n",
    "\n",
    "    # block 1\n",
    "    e1 = self.conv_block(net, filters_root[0])  # 3000*1*64\n",
    "\n",
    "    # block 2\n",
    "    e2 = tf.keras.layers.MaxPool2D(pool_size=(2, 1),padding='same')(e1)  # 1500*1*64\n",
    "    e2 = self.conv_block(e2, filters_root[1])  # 1500*1*128\n",
    "\n",
    "    # block 3\n",
    "    e3 = tf.keras.layers.MaxPool2D(pool_size=(2, 1),padding='same')(e2)  # 750*1*128\n",
    "    e3 = self.conv_block(e3, filters_root[2])  # 750*1*256\n",
    "\n",
    "    # block 4\n",
    "    e4 = tf.keras.layers.MaxPool2D(pool_size=(2, 1),padding='same')(e3)  # 375*1*256\n",
    "    e4 = self.conv_block(e4, filters_root[3])  # 375*1*512\n",
    "\n",
    "    # block 5\n",
    "    # bottleneck layer\n",
    "    e5 =tf.keras.layers.MaxPool2D(pool_size=(3, 1),padding='same')(e4)  # 125*1*512\n",
    "    e5 = self.conv_block(e5, filters_root[4])  # 125*1*1024   \n",
    "\n",
    "    \"\"\" Decoder \"\"\"\n",
    "    cat_channels = filters_root[0]\n",
    "    cat_blocks = len(filters_root)\n",
    "    upsample_channels = cat_blocks * cat_channels\n",
    "\n",
    "    \"\"\" d4 \"\"\"\n",
    "    e1_d4 = tf.keras.layers.MaxPool2D(pool_size=(8, 1),padding='same')(e1)  # e1--3000*1*64 --> 375*1*64\n",
    "    e1_d4 = self.conv_block(e1_d4,cat_channels,n=1)  # 375*1*64 --> 375*1*64\n",
    "\n",
    "    e2_d4 = tf.keras.layers.MaxPool2D(pool_size=(4, 1),padding='same')(e2)  # 1500*1*128 --> 375*40*128\n",
    "    e2_d4 = self.conv_block(e2_d4, cat_channels, n=1)  # 160*160*128 --> 40*40*64\n",
    "\n",
    "    e3_d4 = tf.keras.layers.MaxPool2D(pool_size=(2, 1),padding='same')(e3)  # 750*1*256  --> 375*1*256\n",
    "    e3_d4 = self.conv_block(e3_d4, cat_channels, n=1)  # 375*1*256 --> 375*1*64\n",
    "\n",
    "\n",
    "    e4_d4 = self.conv_block(e4, cat_channels, n=1)  # 375*1*512  --> 375*1*64\n",
    "\n",
    "    e5_d4 = tf.keras.layers.UpSampling2D(size=(3, 1), interpolation='bilinear')(e5)  # e5--125*1*1024  --> 375*1*1024\n",
    "    e5_d4 = self.conv_block(e5_d4, cat_channels, n=1)  #375*1*1024  --> 375*1*64-----cat_channels=64\n",
    "\n",
    "    d4 = tf.keras.layers.concatenate([e1_d4, e2_d4, e3_d4, e4_d4, e5_d4])\n",
    "    d4 = self.conv_block(d4, upsample_channels, n=1)  # 375*1*** --> 375*1*320\n",
    "\n",
    "    \"\"\" d3 \"\"\"\n",
    "    e1_d3 = tf.keras.layers.MaxPool2D(pool_size=(4, 1),padding='same')(e1)  # e1--3000*1*64 --> 750*1*64\n",
    "    e1_d3 = self.conv_block(e1_d3, cat_channels, n=1)  # 750*1*64 --> 750*1*64\n",
    "\n",
    "    e2_d3 = tf.keras.layers.MaxPool2D(pool_size=(2, 1),padding='same')(e2)  #  1500*1*128--> 750*1*128\n",
    "    e2_d3 = self.conv_block(e2_d3, cat_channels, n=1)  # 750*1*128--> 750*1*64\n",
    "\n",
    "    e3_d3 = self.conv_block(e3, cat_channels, n=1)  # 750*1*256 --> 750*1*64\n",
    "\n",
    "    e4_d3 = tf.keras.layers.UpSampling2D(size=(2, 1), interpolation='bilinear')(d4)  # 375*1*320 --> 750*1*320\n",
    "    e4_d3 = self.conv_block(e4_d3, cat_channels, n=1)  # 750*1*320 --> 750*1*64\n",
    "\n",
    "    e5_d3 = tf.keras.layers.UpSampling2D(size=(6, 1), interpolation='bilinear')(e5)  # 125*1*1024  --> 750*1*1024\n",
    "    e5_d3 = self.conv_block(e5_d3, cat_channels, n=1)  # 750*1*1024 --> 750*1*64\n",
    "\n",
    "    d3 = tf.keras.layers.concatenate([e1_d3, e2_d3, e3_d3, e4_d3, e5_d3])\n",
    "    d3 = self.conv_block(d3, upsample_channels, n=1)  # 750*1*** --> 750*1*320\n",
    "\n",
    "    \"\"\" d2 \"\"\"\n",
    "    e1_d2 = tf.keras.layers.MaxPool2D(pool_size=(2, 1),padding='same')(e1)  # 3000*1*64 --> 1500*1*64\n",
    "    e1_d2 = self.conv_block(e1_d2, cat_channels, n=1)  # 1500*1*64 --> 1500*1*64\n",
    "\n",
    "    e2_d2 = self.conv_block(e2, cat_channels, n=1)  # 1500*1*128 --> 1500*1*64\n",
    "\n",
    "    d3_d2 = tf.keras.layers.UpSampling2D(size=(2, 1), interpolation='bilinear')(d3)  # 750*1*320 --> 1500*1*320\n",
    "    d3_d2 = self.conv_block(d3_d2, cat_channels, n=1)  # 1500*1*320--> 1500*1*64\n",
    "\n",
    "    d4_d2 = tf.keras.layers.UpSampling2D(size=(4, 1), interpolation='bilinear')(d4)  # 375*1*320 --> 1500*1*320\n",
    "    d4_d2 = self.conv_block(d4_d2, cat_channels, n=1)  # 1500*1*320 --> 1500*1*64\n",
    "\n",
    "    e5_d2 = tf.keras.layers.UpSampling2D(size=(12, 1), interpolation='bilinear')(e5)  # 125*1*1024 --> 1500*1*1024\n",
    "    e5_d2 = self.conv_block(e5_d2, cat_channels, n=1)  # 1500*1*1024 --> 1500*1*64\n",
    "\n",
    "    d2 = tf.keras.layers.concatenate([e1_d2, e2_d2, d3_d2, d4_d2, e5_d2])\n",
    "    d2 = self.conv_block(d2, upsample_channels, n=1)  # 1500*1*** --> 1500*1*320\n",
    "\n",
    "    \"\"\" d1 \"\"\"\n",
    "    e1_d1 = self.conv_block(e1, cat_channels, n=1)  # 3000*1*64 --> 3000*1*64\n",
    "\n",
    "    d2_d1 = tf.keras.layers.UpSampling2D(size=(2, 1), interpolation='bilinear')(d2)  # 1500*1*320 --> 3000*1*320\n",
    "    d2_d1 = self.conv_block(d2_d1, cat_channels, n=1)  #  3000*1*320 -->  3000*1*64\n",
    "\n",
    "    d3_d1 = tf.keras.layers.UpSampling2D(size=(4, 1), interpolation='bilinear')(d3)  # 750*1*320 --> 3000*1*320\n",
    "    d3_d1 = self.conv_block(d3_d1, cat_channels, n=1)  # 3000*1*320 --> 3000*1*64\n",
    "\n",
    "    d4_d1 = tf.keras.layers.UpSampling2D(size=(8, 1), interpolation='bilinear')(d4)  # 375*1*320 --> 3000*1*320\n",
    "    d4_d1 = self.conv_block(d4_d1, cat_channels, n=1)  # 3000*1*320--> 3000*1*64\n",
    "\n",
    "    e5_d1 = tf.keras.layers.UpSampling2D(size=(24, 1), interpolation='bilinear')(e5)  # 125*1*1024--> 3000*1*1024\n",
    "    e5_d1 = self.conv_block(e5_d1, cat_channels, n=1)  # 3000*1*1024 --> 3000*1*64\n",
    "\n",
    "    d1 = tf.keras.layers.concatenate([e1_d1, d2_d1, d3_d1, d4_d1, e5_d1, ])\n",
    "    d1 = self.conv_block(d1, upsample_channels, n=1)  #3000*1*** --> 3000*1*320\n",
    "\n",
    "    # # 最后一层没有归一化和ReLU\n",
    "\n",
    "    d=tf.keras.layers.Conv2D(self.n_class, kernel_size=(3, 1), padding='same', activation=None)(d1)\n",
    "\n",
    "\n",
    "    output = tf.keras.activations.softsign(d)\n",
    "      \n",
    "    # last layer does not have batch norm and relu\n",
    "    # d1 = self.conv_block(d1, self.n_class, n=1)\n",
    "    # d1 = tf.keras.activations.softmax(d1)\n",
    "\n",
    "    # \"\"\" Deep Supervision Part\"\"\"\n",
    "\n",
    "    # d2=tf.keras.layers.Conv2D(self.n_class,kernel_size=self.kernel_size, padding='same', activation=None)(d2)\n",
    "    # d3=tf.keras.layers.Conv2D(self.n_class,kernel_size=self.kernel_size, padding='same', activation=None)(d3)\n",
    "    # d4=tf.keras.layers.Conv2D(self.n_class,kernel_size=self.kernel_size, padding='same', activation=None)(d4)\n",
    "    # e5=tf.keras.layers.Conv2D(self.n_class,kernel_size=self.kernel_size, padding='same', activation=None)(e5)\n",
    "    \n",
    "\n",
    "    # # d1 = no need for up sampling\n",
    "    # d2 = tf.keras.layers.UpSampling2D(size=(2, 1), interpolation='bilinear')(d2)\n",
    "    # d3 = tf.keras.layers.UpSampling2D(size=(4, 1), interpolation='bilinear')(d3)\n",
    "    # d4 = tf.keras.layers.UpSampling2D(size=(8, 1), interpolation='bilinear')(d4)\n",
    "    # e5 = tf.keras.layers.UpSampling2D(size=(24, 1), interpolation='bilinear')(e5)\n",
    "\n",
    "    # d2 = tf.keras.activations.softmax(d2)\n",
    "    # d3 = tf.keras.activations.softmax(d3)\n",
    "    # d4 = tf.keras.activations.softmax(d4)\n",
    "    # e5 = tf.keras.activations.softmax(e5)\n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "\n",
    "      # 保存编码器部分的最终表示，可以用于其他任务或分析\n",
    "    # with tf.compat.v1.variable_scope(\"representation\"):\n",
    "    #   self.representation = convs[-1]\n",
    "\n",
    "    #logits 是最终的网络输出，未经过激活函数\n",
    "    with tf.compat.v1.variable_scope(\"logits\"):\n",
    "      self.logits = output\n",
    "      # 记录 logits 的直方图摘要，以便在 TensorBoard 中可视化\n",
    "      tmp = tf.compat.v1.summary.histogram(\"logits\", self.logits)\n",
    "      self.summary_train.append(tmp)\n",
    "\n",
    "\n",
    "    # 计算最终预测结果，使用 softmax 将 logits 转换为概率分布\n",
    "    with tf.compat.v1.variable_scope(\"preds\"):\n",
    "      self.preds = tf.nn.softmax(output)\n",
    "      # 记录预测结果的直方图摘要，以便在 TensorBoard 中可视化\n",
    "      tmp = tf.compat.v1.summary.histogram(\"preds\", self.preds)\n",
    "      self.summary_train.append(tmp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "phase",
   "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.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}