image_segmentation_pet_analysis.py

import IPython

import tensorflow as tflow


from tensorflow_examples.models.pix2pix import pix2pix as tflowpix2

import tensorflow_datasets as tflowds

from IPython.display import clear_output
import matplotlib.pyplot as matplot




dataset, info = tflowds.load('oxford_iiit_pet:3.*.*', with_info=True)




def normalize(input_image, input_mask):
  input_image = tflow.cast(input_image, tflow.float32) / 255.0
  input_mask -= 1
  return input_image, input_mask




@tflow.function
def load_image_train(datapoint):
  input_image = tflow.image.resize(datapoint['image'], (128, 128))
  input_mask = tflow.image.resize(datapoint['segmentation_mask'], (128, 128))

  if tflow.random.uniform(()) > 0.5:
    input_image = tflow.image.flip_left_right(input_image)
    input_mask = tflow.image.flip_left_right(input_mask)

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask





def load_image_test(datapoint):
  input_image = tflow.image.resize(datapoint['image'], (128, 128))
  input_mask = tflow.image.resize(datapoint['segmentation_mask'], (128, 128))

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask





TRAIN_LENGTH = info.splits['train'].num_examples
BATCH_SIZE = 64
BUFFER_SIZE = 1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE





train = dataset['train'].map(load_image_train, num_parallel_calls=tflow.data.AUTOTUNE)
test = dataset['test'].map(load_image_test)





train_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()
train_dataset = train_dataset.prefetch(buffer_size=tflow.data.AUTOTUNE)
test_dataset = test.batch(BATCH_SIZE)





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

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

  for i in range(len(display_list)):
    matplot.subplot(1, len(display_list), i+1)
    matplot.title(title[i])
    matplot.imshow(tflow.keras.preprocessing.image.array_to_img(display_list[i]))
    matplot.axis('off')
  matplot.show()





for image, mask in train.take(1):
  sample_image, sample_mask = image, mask
display([sample_image, sample_mask])





OUTPUT_CHANNELS = 3





base_model = tflow.keras.applications.MobileNetV2(input_shape=[128, 128, 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]


down_stack = tflow.keras.Model(inputs=base_model.input, outputs=base_model_outputs)

down_stack.trainable = False





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





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


  skips = down_stack(inputs)
  x = skips[-1]
  skips = reversed(skips[:-1])


  for up, skip in zip(up_stack, skips):
    x = up(x)
    concat = tflow.keras.layers.Concatenate()
    x = concat([x, skip])


  last = tflow.keras.layers.Conv2DTranspose(
      output_channels, 3, strides=2,
      padding='same')  #64x64 -> 128x128

  x = last(x)

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





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




tflow.keras.utils.plot_model(model, show_shapes=True)




def create_mask(pred_mask):
  pred_mask = tflow.argmax(pred_mask, axis=-1)
  pred_mask = pred_mask[..., tflow.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[tflow.newaxis, ...]))])





show_predictions()





class DisplayCallback(tflow.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))




EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_dataset, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=test_dataset,
                          callbacks=[DisplayCallback()])





loss = model_history.history['loss']
val_loss = model_history.history['val_loss']

matplot.figure()
matplot.plot(model_history.epoch, loss, 'r', label='Training loss')
matplot.plot(model_history.epoch, val_loss, 'bo', label='Validation loss')
matplot.title('Training and Validation Loss')
matplot.xlabel('Epoch')
matplot.ylabel('Loss Value')
matplot.ylim([0, 1])
matplot.legend()
matplot.show()





show_predictions(test_dataset, 3)





try:
  model_history = model.fit(train_dataset, epochs=EPOCHS,
                            steps_per_epoch=STEPS_PER_EPOCH,
                            class_weight = {0:2.0, 1:2.0, 2:1.0})
  assert False
except Exception as e:
  print(f"{type(e).__name__}: {e}")





label = [0,0]
prediction = [[-3., 0], [-3, 0]] 
sample_weight = [1, 10] 

loss = tflow.losses.SparseCategoricalCrossentropy(from_logits=True,
                                               reduction=tflow.losses.Reduction.NONE)
loss(label, prediction, sample_weight).numpy()





def add_sample_weights(image, label):

  class_weights = tflow.constant([2.0, 2.0, 1.0])
  class_weights = class_weights/tflow.reduce_sum(class_weights)


  sample_weights = tflow.gather(class_weights, indices=tflow.cast(label, tflow.int32))

  return image, label, sample_weights





train_dataset.map(add_sample_weights).element_spec





weighted_model = unet_model(OUTPUT_CHANNELS)
weighted_model.compile(
    optimizer='adam',
    loss=tflow.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'])




weighted_model.fit(
    train_dataset.map(add_sample_weights),
    epochs=1,
    steps_per_epoch=10)