# Databricks notebook source
# MAGIC %md
# MAGIC ScaDaMaLe Course [site](https://lamastex.github.io/scalable-data-science/sds/3/x/) and [book](https://lamastex.github.io/ScaDaMaLe/index.html)

# COMMAND ----------

# MAGIC %md
# MAGIC # U-Net model for image segmentation
# MAGIC 
# MAGIC This is a modified version of Tensorflows tutorial regarding image segmentation which can be found [here](https://www.tensorflow.org/tutorials/images/segmentation). Using a modified [U-Net](https://arxiv.org/abs/1505.04597) approach, with a [VGG16](https://arxiv.org/abs/1409.1556) as the encoder and then using traditional Conv2DTranspose layers for upsampling the dimensions. After 1 epoch a validation accuracy of 84.5 % was achieved on the [Oxford Pets Data Set](https://www.robots.ox.ac.uk/~vgg/data/pets/).

# COMMAND ----------

import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
from IPython.display import clear_output

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

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