{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Building U-Net from scratch for image segmentation","metadata":{}},{"cell_type":"markdown","source":"# What is Image Segmentation\nmage segmentation is a method in which a digital image is broken down into various subgroups called Image segments which helps in reducing the complexity of the image to make further processing or analysis of the image simpler. \nSegmentation in easy words is assigning labels to pixels. All picture elements or pixels belonging to the same category have a common label assigned to them.  \nImage segmentation is a function that takes image inputs and produces an output. The output is a mask or a matrix with various elements specifying the object class or instance to which each pixel belongs.  \nTraditional image segmentation techniques based on such heuristics can be fast and simple, but they often require significant fine-tuning to support specific use cases with manually designed heuristics. They are not always sufficiently accurate to use for complex images. Newer segmentation techniques use machine learning and deep learning to increase accuracy and flexibility.  \nMachine learning-based image segmentation approaches use model training to improve the program’s ability to identify important features. Deep neural network technology is especially effective for image segmentation tasks. ","metadata":{}},{"cell_type":"markdown","source":"# What is U-Net\nThe U-Net architecture was proposed in the [U-Net: Convolutional Networks for Biomedical Image Segmentation paper](https://arxiv.org/pdf/1505.04597.pdf) in 2015.  \nWith this U-Net architecture, the segmentation of images of sizes 512X512 can be computed with a modern GPU within small amounts of time. There have been many variants and modifications of this architecture due to its phenomenal success. Some of them include LadderNet, U-Net with attention, the recurrent and residual convolutional U-Net (R2-UNet), and U-Net with residual blocks or blocks with dense connections.  \n**Building blocks of U-Net :**\n- Encoder : series of layers that extract image features using progressively deeper, narrower filters. The encoder might be pre-trained on a similar task (e.g., image recognition), allowing it to leverage its existing knowledge to perform segmentation tasks.\n- Decoder : series of layers that gradually convert the encoder’s output into a segmentation mask corresponding with the input image’s pixel resolution.\n- Skip connections : multiple long-range neural network connections allowing the model to identify features at different scales to enhance model accuracy.\n\n**Architecture of the U-Net**  \n<p align=\"center\">\n  <img src=\"https://blog.paperspace.com/content/images/size/w1000/2021/05/image-26.png\" />\n</p>\nThe contracting and expansive paths mirror each other. The name U-Net comes from the shape of the network, which forms a U-shape.  \nThe encoder block has a constant reduction of image size with the help of the max-pooling layers of strides 2. We also have repeated convolutional layers with an increasing number of filters in the encoder architecture. Once we reach the decoder aspect, we notice the number of filters in the convolutional layers start to decrease along with a gradual upsampling in the following layers all the way to the top. We also notice that the use of skip connections that connect the previous outputs with the layers in the decoder blocks.  \nThis skip connection is a vital concept to preserve the loss from the previous layers so that they reflect stronger on the overall values. High-resolution features from the contracting path are combined with the upsampled output to improve the network's localization capacity.","metadata":{}},{"cell_type":"markdown","source":"## Implementing U-Net architecture\nAs an example, I'll be segmenting pets from the recognised [Oxford Pets dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/). I'll build and utilise the U-Net to perform image segmentation.  \nThe Oxford pet dataset consists of 37 categories of pet dataset with roughly 200 images for each class. The images have a large variation in scale, pose, and lighting.","metadata":{}},{"cell_type":"markdown","source":"### Downloading the data\nThe data is available at the above highlighted link. Although for this, I have attached the data along with the notebook, thanks to Kaggle.","metadata":{}},{"cell_type":"markdown","source":"## Preparing the paths of input images and target segmentation masks\nAlong with setting hyper-parameters such as image size, batch size and number of classes, we'll set up and ensure that the elements in the dataset are arranged in the correct order for performing and obtaining a satisfactory image segmentation task","metadata":{}},{"cell_type":"code","source":"import os\n\ninput_dir = \"/kaggle/input/the-oxfordiiit-pet-dataset/images/images/\"\ntarget_dir = \"/kaggle/input/the-oxfordiiit-pet-dataset/annotations/annotations/trimaps\"\nimg_size = (160, 160)\nnum_classes = 3\nbatch_size = 32\n\ninput_img_paths = sorted(\n    [\n        os.path.join(input_dir, fname)\n        for fname in os.listdir(input_dir)\n        if fname.endswith(\".jpg\")\n    ]\n)\ntarget_img_paths = sorted(\n    [\n        os.path.join(target_dir, fname)\n        for fname in os.listdir(target_dir)\n        if fname.endswith(\".png\") and not fname.startswith(\".\")\n    ]\n)\n\nprint(\"Number of samples:\", len(input_img_paths))\n\nfor input_path, target_path in zip(input_img_paths[:10], target_img_paths[:10]):\n    print(input_path, \"|\", target_path)","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:08.662633Z","iopub.execute_input":"2022-12-19T16:59:08.663524Z","iopub.status.idle":"2022-12-19T16:59:10.151291Z","shell.execute_reply.started":"2022-12-19T16:59:08.663405Z","shell.execute_reply":"2022-12-19T16:59:10.150230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Exploring and Visualizing the data\nOur next step will be to take a brief look at the dataset. Let's analyze the dataset by displaying both an image and its respective segmented output. This segmented output with the masking is often times referred to as the **ground truth annotation**.","metadata":{}},{"cell_type":"code","source":"from IPython.display import Image, display\nfrom tensorflow.keras.preprocessing.image import load_img\nimport PIL\nfrom PIL import ImageOps\n\n# Display input image #2\ndisplay(Image(filename=input_img_paths[2]))\n\n# Display auto-contrast version of corresponding target (per-pixel categories)\nimg = PIL.ImageOps.autocontrast(load_img(target_img_paths[2]))\ndisplay(img)","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:10.155954Z","iopub.execute_input":"2022-12-19T16:59:10.156827Z","iopub.status.idle":"2022-12-19T16:59:15.543884Z","shell.execute_reply.started":"2022-12-19T16:59:10.156788Z","shell.execute_reply":"2022-12-19T16:59:15.542927Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Configuring data generator\nWe will utilize the Sequence operation from the Keras deep learning framework. This method is much safer for multi-processing in comparison to using the generator because they only train each sample once per epoch, which will avoid any unnecessary adjustments that we will need to make.","metadata":{}},{"cell_type":"code","source":"from tensorflow.keras import utils\nimport numpy as np\nfrom tensorflow.keras.preprocessing.image import load_img\n\n\nclass OxfordPetsDataGenerator(utils.Sequence):\n    \"\"\"Helper to iterate over the data (as Numpy arrays).\"\"\"\n\n    def __init__(self, batch_size, img_size, input_img_paths, target_img_paths):\n        self.batch_size = batch_size\n        self.img_size = img_size\n        self.input_img_paths = input_img_paths\n        self.target_img_paths = target_img_paths\n\n    def __len__(self):\n        return len(self.target_img_paths) // self.batch_size\n\n    def __getitem__(self, idx):\n        \"\"\"Returns tuple (input, target) correspond to batch #idx.\"\"\"\n        i = idx * self.batch_size\n        batch_input_img_paths = self.input_img_paths[i : i + self.batch_size]\n        batch_target_img_paths = self.target_img_paths[i : i + self.batch_size]\n        \n        # Filling in the feature image data\n        x = np.zeros((self.batch_size,) + self.img_size + (3,), dtype=\"float32\")\n        for j, path in enumerate(batch_input_img_paths):\n            \n            # loads in the image and resizes in the target size\n            img = load_img(path, target_size=self.img_size)\n            x[j] = img\n            \n        # Filling in the label data\n        y = np.zeros((self.batch_size,) + self.img_size + (1,), dtype=\"uint8\")\n        for j, path in enumerate(batch_target_img_paths):\n            img = load_img(path, target_size=self.img_size, color_mode=\"grayscale\")\n            y[j] = np.expand_dims(img, 2)\n            # Ground truth labels are 1, 2, 3. Subtract one to make them 0, 1, 2:\n            y[j] -= 1\n        \n        return x, y","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:15.545516Z","iopub.execute_input":"2022-12-19T16:59:15.546414Z","iopub.status.idle":"2022-12-19T16:59:15.559374Z","shell.execute_reply.started":"2022-12-19T16:59:15.546373Z","shell.execute_reply":"2022-12-19T16:59:15.555464Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Splitting the data into training and validation sets\nBoth of these data entities must be viewed separately so that the model does not get a peek at the testing data. In the validation set, we will also perform an optional shuffle operation that will mix up all the images in the dataset, and we can obtain random samples for both the train and the validation images.","metadata":{}},{"cell_type":"code","source":"import random\n\n# Split our img paths into a training and a validation set\nval_samples = 1000\nrandom.Random(42).shuffle(input_img_paths)\nrandom.Random(42).shuffle(target_img_paths)\n\n# Splitting data accoring to val_samples\ntrain_input_img_paths = input_img_paths[:-val_samples]\ntrain_target_img_paths = target_img_paths[:-val_samples]\n\nval_input_img_paths = input_img_paths[-val_samples:]\nval_target_img_paths = target_img_paths[-val_samples:]\n\n# Instantiating data generators for training and validation sets\ntrain_gen = OxfordPetsDataGenerator(batch_size, img_size, train_input_img_paths, train_target_img_paths)\nval_gen = OxfordPetsDataGenerator(batch_size, img_size, val_input_img_paths, val_target_img_paths)","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:15.562363Z","iopub.execute_input":"2022-12-19T16:59:15.562780Z","iopub.status.idle":"2022-12-19T16:59:15.590279Z","shell.execute_reply.started":"2022-12-19T16:59:15.562735Z","shell.execute_reply":"2022-12-19T16:59:15.589388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Creating the U-Net architecture\nWe'll build the convolutional, encoder and decoder separately and then assemble these for better ease in making the model.","metadata":{}},{"cell_type":"code","source":"# Importing libraries and modules\nimport warnings\nimport tensorflow as tf \nfrom tensorflow import keras\nfrom tensorflow.keras import layers, models\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\nwarnings.filterwarnings('ignore')\n\n# Building the convolutional block\ndef ConvBlock(inputs, filters=64):\n    # Taking first input and implementing the first conv block\n    conv1 = layers.Conv2D(filters, kernel_size = (3,3), padding = \"same\")(inputs)\n    batch_norm1 = layers.BatchNormalization()(conv1)\n    act1 = layers.ReLU()(batch_norm1)\n    \n    # Taking first input and implementing the second conv block\n    conv2 = layers.Conv2D(filters, kernel_size = (3,3), padding = \"same\")(act1)\n    batch_norm2 = layers.BatchNormalization()(conv2)\n    act2 = layers.ReLU()(batch_norm2)\n    \n    return act2\n\n# Building the encoder\ndef encoder(inputs, filters=64):\n    # Collect the start and end of each sub-block for normal pass and skip connections\n    enc1 = ConvBlock(inputs, filters)\n    MaxPool1 = layers.MaxPooling2D(strides = (2,2))(enc1)\n    return enc1, MaxPool1\n\n# Building the decoder\ndef decoder(inputs, skip, filters=64):\n    # Upsampling and concatenating the essential features\n    Upsample = layers.Conv2DTranspose(filters, (2, 2), strides=2, padding=\"same\")(inputs)\n    Connect_Skip = layers.Concatenate()([Upsample, skip])\n    out = ConvBlock(Connect_Skip, filters)\n    return out\n\ndef get_U_Net(image_size):\n    inputs = layers.Input(image_size)\n    \n    # Construct the encoder blocks and increasing the filters by a factor of 2\n    skip1, encoder_1 = encoder(inputs, 64)\n    skip2, encoder_2 = encoder(encoder_1, 64*2)\n    skip3, encoder_3 = encoder(encoder_2, 64*4)\n    skip4, encoder_4 = encoder(encoder_3, 64*8)\n    \n    # Preparing the next block\n    conv_block = ConvBlock(encoder_4, 64*16)\n    \n    # Construct the decoder blocks and decreasing the filters by a factor of 2\n    decoder_1 = decoder(conv_block, skip4, 64*8)\n    decoder_2 = decoder(decoder_1, skip3, 64*4)\n    decoder_3 = decoder(decoder_2, skip2, 64*2)\n    decoder_4 = decoder(decoder_3, skip1, 64)\n    \n    outputs = layers.Conv2D(3, 1, padding=\"same\", activation=\"sigmoid\")(decoder_4)\n    \n    model = models.Model(inputs, outputs)\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:15.592373Z","iopub.execute_input":"2022-12-19T16:59:15.592687Z","iopub.status.idle":"2022-12-19T16:59:15.608062Z","shell.execute_reply.started":"2022-12-19T16:59:15.592662Z","shell.execute_reply":"2022-12-19T16:59:15.607009Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Finalizing the model","metadata":{}},{"cell_type":"code","source":"input_shape = (160, 160, 3)\nmodel = get_U_Net(input_shape)\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:15.610336Z","iopub.execute_input":"2022-12-19T16:59:15.611084Z","iopub.status.idle":"2022-12-19T16:59:18.851429Z","shell.execute_reply.started":"2022-12-19T16:59:15.611048Z","shell.execute_reply":"2022-12-19T16:59:18.850428Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.keras.utils.plot_model(model, \"model.png\", show_shapes=False, show_dtype=False, show_layer_names=True, rankdir='TB', expand_nested=False, dpi=96)","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:18.853049Z","iopub.execute_input":"2022-12-19T16:59:18.853732Z","iopub.status.idle":"2022-12-19T16:59:20.034177Z","shell.execute_reply.started":"2022-12-19T16:59:18.853693Z","shell.execute_reply":"2022-12-19T16:59:20.031310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training the model\nWe will now compile and train the model to see its performance on the data.","metadata":{}},{"cell_type":"code","source":"model.compile(optimizer=\"Adam\", loss=\"sparse_categorical_crossentropy\")\n\ncallbacks = [\n    keras.callbacks.ModelCheckpoint(\"oxford_segmentation.h5\", save_best_only=True)\n]\n\nepochs = 15","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:39.760346Z","iopub.execute_input":"2022-12-19T16:59:39.761296Z","iopub.status.idle":"2022-12-19T16:59:39.779822Z","shell.execute_reply.started":"2022-12-19T16:59:39.761255Z","shell.execute_reply":"2022-12-19T16:59:39.778898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train the model, doing validation at the end of each epoch.\nmodel.fit(train_gen, epochs=epochs, validation_data=val_gen, callbacks=callbacks)","metadata":{"execution":{"iopub.status.busy":"2022-12-19T16:59:40.038643Z","iopub.execute_input":"2022-12-19T16:59:40.038993Z","iopub.status.idle":"2022-12-19T17:20:26.589352Z","shell.execute_reply.started":"2022-12-19T16:59:40.038964Z","shell.execute_reply":"2022-12-19T17:20:26.588413Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Viewing the predictions of the constructed U-Net model","metadata":{}},{"cell_type":"code","source":"# Generate predictions for all images in the validation set\n\nval_gen = OxfordPetsDataGenerator(batch_size, img_size, val_input_img_paths, val_target_img_paths)\nval_preds = model.predict(val_gen)\n\n\ndef display_mask(i):\n    \"\"\"Quick utility to display a model's prediction.\"\"\"\n    mask = np.argmax(val_preds[i], axis=-1)\n    mask = np.expand_dims(mask, axis=-1)\n    img = PIL.ImageOps.autocontrast(keras.preprocessing.image.array_to_img(mask))\n    display(img)\n\n\n# Display results for validation image idx 2\ni = 2\n\n# Display input image\ndisplay(Image(filename=val_input_img_paths[i]))\n\n# Display ground-truth target mask\nimg = PIL.ImageOps.autocontrast(load_img(val_target_img_paths[i]))\ndisplay(img)\n\n# Display mask predicted by our model\ndisplay_mask(i)  # Note that the model only sees inputs at 150x150.","metadata":{"execution":{"iopub.status.busy":"2022-12-19T17:20:30.923975Z","iopub.execute_input":"2022-12-19T17:20:30.924330Z","iopub.status.idle":"2022-12-19T17:20:39.114144Z","shell.execute_reply.started":"2022-12-19T17:20:30.924298Z","shell.execute_reply":"2022-12-19T17:20:39.113032Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Conclusion\nThe U-Net architecture is one of the most significant and revolutionary landmarks in the field of deep learning. While the initial research paper that introduced the U-Net architecture was to solve the task of Biomedical Image Segmentation, it was not limited to this single application. The model could and can still solve the most complex problems in deep learning. Although some of the elements in the original architecture are outdated, there are several variations of this architecture. These include LadderNet, U-Net with attention, the recurrent and residual convolutional U-Net (R2-UNet), and other similar networks which are derived successfully from the original U-Net Models.","metadata":{}}]}