code: def unet_model(output_channels):
      # This is the last layer of the model
      last = tf.keras.layers.Conv2DTranspose(
          output_channels,
          3,
          strides=2,
          padding="same",
          activation="softmax",
      )  # 64x64 -> 128x128
  
      inputs = tf.keras.layers.Input(shape=[128, 128, 3])
      x = inputs
  
      # Downsampling through the model
      skips = down_stack(x)
      x = skips[-1]
      skips = reversed(skips[:-1])
  
      # Upsampling and establishing the skip connections
      for up, skip in zip(up_stack, skips):
          x = up(x)
          concat = tf.keras.layers.Concatenate()
          x = concat([x, skip])
  
      x = last(x)
  
      return tf.keras.Model(inputs=inputs, outputs=x)
  
file_path: TensorFlowOnSpark/examples/segmentation/segmentation.py
imports: from __future__ import absolute_import, division, print_function, unicode_literals
  from tensorflow_examples.models.pix2pix import pix2pix
  import tensorflow_datasets
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: unet_model(output_channels) and returns: tf