<jupyter_start><jupyter_text><jupyter_code>import numpy as np
import tensorflow as tf
import matplotlib as plt
from tensorflow import keras
from tensorflow.keras.layers import Conv2D, concatenate, Dropout,MaxPool2D, MaxPooling2D, Conv2DTranspose, Activation, BatchNormalization,UpSampling2D, Add

a  = np.random.randn(1,16,16,18)
b =  tf.keras.layers.UpSampling2D(size=(8, 8), interpolation="nearest")(a)
b.shape
c =np.random.randn(1,128,128,18)
d = tf.keras.layers.concatenate([b,c])
d.shape
def conv_2d_block(x,n_filters,k_size,batchnorm=False):
  ''' add two Conv layers with relu activation '''
  #first layer
  x = Conv2D(filters=n_filters,kernel_size=(k_size,k_size) , 
             padding='same', kernel_initializer = 'he_normal')(x)
  if batchnorm:
    x = BatchNormalization()(x)
  x = Activation('relu')(x)
  # 2nd layer
  x = Conv2D(filters=n_filters,kernel_size=(k_size,k_size) , 
              padding='same', kernel_initializer = 'he_normal')(x)
  if batchnorm:
      x = BatchNormalization()(x)
  x = Activation('relu')(x)
  
  x = Conv2D(filters=n_filters,kernel_size=(k_size,k_size) , 
              padding='same', kernel_initializer = 'he_normal')(x)
  if batchnorm:
      x = BatchNormalization()(x)
  x = Activation('relu')(x)

  return x

def u_net(input,n_filters=16,conv_k_size=3,pool_size=2,batchnorm=True,dropout=.2):

  c1 = conv_2d_block(input,n_filters * 1 , conv_k_size,batchnorm)
  p1 = MaxPool2D(pool_size=(pool_size,pool_size))(c1)
  p1 = Dropout(dropout)(p1)

  c2 = conv_2d_block(p1,n_filters * 2 , conv_k_size,batchnorm)
  p2 = MaxPool2D(pool_size=(pool_size,pool_size))(c2)
  p2 = Dropout(dropout)(p2)

  c3 = conv_2d_block(p2,n_filters * 4 , conv_k_size,batchnorm)
  p3 = MaxPool2D(pool_size=(pool_size,pool_size))(c3)
  p3 = Dropout(dropout)(p3)
 
  c4 = conv_2d_block(p3,n_filters * 8 , conv_k_size,batchnorm)
  p4 = MaxPool2D(pool_size=(pool_size,pool_size))(c4)
  p4 = Dropout(dropout)(p4)

  c5 = conv_2d_block(p4,n_filters * 16 , conv_k_size,batchnorm)

   #Up sampling
  u6 = Conv2DTranspose(filters=n_filters * 8 ,kernel_size=(3,3), strides=(2,2),padding='same')(c5)
  u6 = concatenate([u6,c4])
  u6 = Dropout(dropout)(u6)
  c7 = conv_2d_block(u6,n_filters * 8 , conv_k_size,batchnorm)

  u8 = Conv2DTranspose(filters=n_filters * 4 ,kernel_size=(3,3), strides=(2,2),padding='same')(c7)
  u8 = concatenate([u8,c3])
  u8 = Dropout(dropout)(u8)
  c9 = conv_2d_block(u8,n_filters * 4 , conv_k_size,batchnorm)

  u10 = Conv2DTranspose(filters=n_filters * 2,kernel_size=(3,3) , strides=(2,2),padding='same')(c9)
  u10 = concatenate([u10,c2])
  u10 = Dropout(dropout)(u10)
  c11 = conv_2d_block(u10,n_filters * 2 , conv_k_size,batchnorm)

  u12 = Conv2DTranspose(filters=n_filters * 1 ,kernel_size=(3,3), strides=(2,2),padding='same')(c11)
  u12 = concatenate([u12,c1])
  u12 = Dropout(dropout)(u12)
  c13 = conv_2d_block(u12,n_filters * 1 , conv_k_size,batchnorm)

  output = Conv2D(filters=3 , kernel_size=(1,1),activation='softmax')(c13) # output layer
 # model = Model(inputs=input,outputs=output,name='classifier')
  return output
u_net(d).shape<jupyter_output><empty_output>