Tensorflow failed segmentation

import numpy as np
import os

import imageio
import PIL
import PIL.Image as Image

import matplotlib as mpl
import matplotlib.pyplot as plt
from IPython.display import clear_output

import tensorflow as tf
import tensorflow_datasets as tfds
from examples.tensorflow_examples.models.pix2pix import pix2pix
from sklearn.model_selection import train_test_split

########################################################################
### DEFINES

HEIGHT = 128
WIDTH = 128
BATCH_SIZE = 10

SEED = 154

DIR_DAT = "" // datasetfolder/data
DIR_MASK = "" // datasetfolder/mask


TRAIN_LENGTH = 50
EPOCHS = 20
VAL_SUBSPLITS = 1

OUTPUT_CLASSES = 3

# Define a normalizer that can be applied while visualizing masks to have a consistency
NORM = mpl.colors.Normalize(vmin=0, vmax=58)

########################################################################
### FUNCTIONS

def get_paths(image_dir_path, mask_dir_path):
  # Walk through both paths and input the concatenated string for every file
  image_path = []
  for root, dirs, files in os.walk(image_dir_path):
      for file in files:
          path = os.path.join(root,file)
          image_path.append(path)

  mask_path = []
  for root, dirs, files in os.walk(mask_dir_path):
      for file in files:
          path = os.path.join(root,file)
          mask_path.append(path)

  # The one time this is desirable
  image_path.sort()
  mask_path.sort()
  return image_path, mask_path

def get_images(image_path, mask_path):
  # create a list to store images
  images = []
  for path in image_path:
      # read file
      file = tf.io.read_file(path)
      # decode png file into a tensor
      image = tf.image.decode_png(file, channels=3)
      # append to the list
      images.append(image)

  # create a list to store masks
  masks = []
  for path in mask_path:
      # read the file
      file = tf.io.read_file(path)
      # decode png file into a tensor
      mask = tf.image.decode_png(file, channels=1)
      # append mask to the list
      masks.append(mask)
  return images, masks

def show_images(images):
  plt.figure(figsize=(48,32))

  for i in range(1,7):
      # Create a subplot for each image
      plt.subplot(2,6,i)
      # Get the i-th image from the list
      img = images[i]
      # Show the image with a colorbar
      plt.imshow(img)
      plt.colorbar()
      # Turn off the axis labels
      plt.axis('off')

def show_masks(masks):
  for i in range(1,7):
      plt.subplot(2,6,6+i)
      img = masks[i]
      plt.imshow(img, cmap='jet', norm=NORM)
      plt.colorbar()
      plt.axis('off')
  plt.show()

def resize_image(image):
    # scale the image
    image = tf.cast(image, tf.float32)
    image = image/255.0
    # resize image
    image = tf.image.resize(image, (128,128))
    return image

def resize_mask(mask):
    # resize the mask
    mask = tf.image.resize(mask, (128,128))
    mask = tf.cast(mask, tf.uint8)
    return mask

def unet_model(output_channels:int):
  inputs = tf.keras.layers.Input(shape=[128, 128, 3])

  # Downsampling through the model
  skips = down_stack(inputs)
  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])

  # This is the last layer of the model
  last = tf.keras.layers.Conv2DTranspose(
      filters=output_channels, kernel_size=3, strides=2,
      padding='same')  #64x64 -> 128x128

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)

def display(display_list):
  plt.figure(figsize=(15, 15))

  title = ['Input Image', 'True Mask', 'Predicted Mask']

  for i in range(len(display_list)):
    plt.subplot(1, len(display_list), i+1)
    plt.title(title[i])
    plt.imshow(tf.keras.utils.array_to_img(display_list[i]))
    plt.axis('off')
  plt.show()

def create_mask(pred_mask):
  pred_mask = tf.math.argmax(pred_mask, axis=-1)
  pred_mask = pred_mask[..., tf.newaxis]
  return pred_mask[0]

def show_predictions(dataset=None, num=1):
  if dataset:
    for image, mask in dataset.take(num):
      pred_mask = model.predict(image)
      display([image[0], mask[0], create_mask(pred_mask)])
  else:
    display([sample_image, sample_mask,
             create_mask(model.predict(sample_image[tf.newaxis, ...]))])

########################################################################
### CODE

paths_im, paths_msk = get_paths(DIR_DAT, DIR_MASK)
images, masks = get_images(paths_im, paths_msk)
paths_im.clear()
paths_msk.clear()

#show_images(images)
#show_masks(masks)

images_pre = [resize_image(i) for i in images]
masks_pre = [resize_mask(m) for m in masks]

sample_image = images_pre[0]
sample_mask = masks_pre[0]

#show_images(images_pre)
#show_masks(masks_pre)

train_X, val_X,train_y, val_y = train_test_split(images_pre, masks_pre, test_size=0.2,
                                                      random_state=0
                                                     )

# develop tf Dataset objects
train_X = tf.data.Dataset.from_tensor_slices(train_X)
val_X = tf.data.Dataset.from_tensor_slices(val_X)

train_y = tf.data.Dataset.from_tensor_slices(train_y)
val_y = tf.data.Dataset.from_tensor_slices(val_y)

# Verify shapes
#print(train_X.element_spec)
#print(train_y.element_spec)
#print(val_X.element_spec)
#print(val_y.element_spec)

train = tf.data.Dataset.zip((train_X, train_y))
val = tf.data.Dataset.zip((val_X, val_y))

train = train.batch(BATCH_SIZE, drop_remainder=True)
print(train.element_spec)
val = val.batch(BATCH_SIZE, drop_remainder=True)
print(val.element_spec)
train = train.repeat(EPOCHS)


### Base model

base_model = tf.keras.applications.MobileNetV2(input_shape=[WIDTH, HEIGHT, 3], include_top=False)
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
down_stack.trainable = False

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]

model = unet_model(output_channels=OUTPUT_CLASSES)
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# Isn't she beautiful
tf.keras.utils.plot_model(model, show_shapes=True, expand_nested=True, dpi=64)

show_predictions()

class DisplayCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    clear_output(wait=True)
    show_predictions()
    print ('\nSample Prediction after epoch {}\n'.format(epoch+1))

STEPS_PER_EPOCH = len(train)
VALIDATION_STEPS = len(val)

model_history = model.fit(train, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=val)

show_predictions()