{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "luSS6e7lSL0H" }, "source": [ "# Module 2: U-Net for Cell Segmentation\n", "---\n", "\n", "Welcome to the practical session of [Data-Driven Life Sciences course module 4](https://ddls.aicell.io/course/ddls-2025/module-2/lab/), created by [Estibaliz Gómez de Mariscal](https://esgomezm.github.io/). And, Professor Wei Ouyang, teaching assistants Songtao Cheng, and Nils Mechtel made some modifications.\n", "\n", "## Introduction\n", "This is a notebook that shows how to design and train a [U-Net](https://en.wikipedia.org/wiki/U-Net)-like network to segment cells in Phase Contrast Microscopy images using Keras and Tensorflow. The aim is to train the network using original phase contrast microscopy images as input, and one label image per category (background, foreground and contours) as output.\n", "\n", "
\n", "
\n", "\n", "
\n", "\n", "These notebook is based on the previous development of the following authors for the [NEUBIAS text book 2021](https://github.com/NEUBIAS/neubias-springer-book-2021).\n", "\n", "**Authors**: **Authors**: [Estibaliz Gómez-de-Mariscal](https://henriqueslab.github.io/team/2021-10-01-EGdM/), [Daniel Franco-Barranco](https://danifranco.github.io), [Arrate Muñoz-Barrutia](https://image.hggm.es/es/arrate-munoz) and [Ignacio Arganda-Carreras](https://sites.google.com/site/iargandacarreras/).\n", "\n", "## Data\n", "The image data used in the notebook was provided by Dr. T. Becker. Fraunhofer Institution for Marine Biotechnology, Lübeck, Germany to the [Cell Tracking Challenge](http://celltrackingchallenge.net/). The acquisition details can be found [here](http://celltrackingchallenge.net/2d-datasets/#bg-showmore-hidden-5ea69464435099020168051).\n", "The training data set consists of 2D images of pancreatic stem cells on a polystyrene substrate. The field of view measures 1152 x 922 microns approx., with a resolution of 1.6 x 1.6 um/pixel. In this notebook we will use the gold reference tracking annotations (ST) of the training data (frames 150-250 of both videos). The annotations were binarized for this use-case. Find the data adapted for this notebook [here](https://github.com/esgomezm/NEUBIAS_chapter_DL_2020/releases/download/1/data4notebooks.zip):\n", "\n", "Frames of sequence 01 were used as training data and frames of the sequence 02 were used validation (frames 150, 140, 150, ..., 250) and test data (frames 151, 152, ..., 248, 249).\n", "\n", "The data is organized in different folders as follows:\n", "\n", "```\n", "./\n", " |-- train_input\n", " | | t150.tif\n", " | | ...\n", " |-- train_binary_masks\n", " | | man_seg150.tif\n", " | | ...\n", " |-- train_contours\n", " | | man_seg150.png\n", " | | ...\n", " |-- validation_input\n", " | | t150.tif\n", " | | ...\n", " |-- validation_binary_masks\n", " | | man_seg150.tif\n", " | | ...\n", " |-- validation_contours\n", " | | man_seg150.png\n", " | | ...\n", " |-- test_input\n", " | | t151.tif\n", " | | ...\n", " |-- test_binary_masks\n", " | | man_seg151.tif\n", " | | ...\n", " |-- test_contours\n", " | | man_seg151.png\n", " | | ...\n", "```\n", "\n", "\n", "## Network arquitecture:\n", "\n", "We will train the standard encoder-decoder called U-Net.\n", "
\n", "
\n", "\n", "
\n", "\n", "This network was introduced by\n", "\n", "- U-Net: Convolutional Networks for Biomedical Image Segmentation by Ronneberger et al. published on arXiv in 2015 (https://arxiv.org/abs/1505.04597)\n", "\n", "and\n", "\n", "- U-Net: deep learning for cell counting, detection, and morphometry by Thorsten Falk et al. in Nature Methods 2019 (https://www.nature.com/articles/s41592-018-0261-2) And source code found in: https://github.com/zhixuhao/unet by Zhixuhao" ] }, { "cell_type": "markdown", "metadata": { "id": "iSckjyNSp6D9" }, "source": [ "## Getting started\n", "\n", "### Familiarize yourself with the topic:\n", "\n", "- Ensure you're comfortable with bioimage analysis, data augmentation, and U-Nets. You can create a personalized prompt in ChatGPT to help guide you through these topics.\n", "- Remember to upload your ChatGPT conversation history in the submission form.\n", "\n", "### Important Note for This Lab Notebook:\n", "\n", "- **🌞 Tasks Introduction:** Sections marked with a 🌞 symbol introduce an exercise or question. Please read these sections carefully to understand the concepts and tasks involved.\n", "\n", "- **⭐ Your Answer Here:** Cells marked with a ⭐ symbol indicate where you need to write your answer. Please provide your code or answer there.\n", "\n", "### Now, let's start:\n" ] }, { "cell_type": "markdown", "metadata": { "id": "itnfhtxN7zVE" }, "source": [ "### Download the ZIP file with the image data and unzip it in Google's content." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "xcJn0lZtEq9z" }, "outputs": [], "source": [ "import zipfile\n", "\n", "# Download file with image data\n", "!wget 'https://github.com/esgomezm/NEUBIAS_chapter_DL_2020/releases/download/1/data4notebooks.zip'\n", "path2zip= 'data4notebooks.zip'\n", "\n", "# Extract locally\n", "with zipfile.ZipFile(path2zip, 'r') as zip_ref:\n", " zip_ref.extractall('dataset/')" ] }, { "cell_type": "markdown", "metadata": { "id": "SK-GgwqbrkI2" }, "source": [ "We should be able to read the list of **101 training images**, together with their corresponding binary masks and cell contours." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "aIuhK8x0-wBC" }, "outputs": [], "source": [ "import os\n", "# Path to the training images\n", "train_input_path = 'dataset/train_input'\n", "train_masks_path = 'dataset/train_binary_masks'\n", "train_contours_path = 'dataset/train_contours'\n", "# Read the list of file names and sort them to have a match between images and masks\n", "train_input_filenames = [x for x in os.listdir( train_input_path ) if x.endswith(\".tif\")]\n", "train_input_filenames.sort()\n", "train_masks_filenames = [x for x in os.listdir( train_masks_path ) if x.endswith(\".tif\")]\n", "train_masks_filenames.sort()\n", "train_contours_filenames = [x for x in os.listdir( train_contours_path ) if x.endswith(\".png\")]\n", "train_contours_filenames.sort()\n", "\n", "print( 'Number of training input images: ' + str( len(train_input_filenames)) )\n", "print( 'Number of training binary mask images: ' + str( len(train_masks_filenames)) )\n", "print( 'Number of training contour images: ' + str( len(train_contours_filenames)) )" ] }, { "cell_type": "markdown", "metadata": { "id": "nDfRdihFrw74" }, "source": [ "Next, we read all those images into memory and display one with its corresponding labels (masks and countours)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "hckMCAdLEn8W" }, "outputs": [], "source": [ "import cv2\n", "from matplotlib import pyplot as plt\n", "\n", "# Read training images (input, mask and contours)\n", "train_img = [cv2.imread(os.path.join(train_input_path, x),\n", " cv2.IMREAD_ANYDEPTH) for x in train_input_filenames ]\n", "train_masks = [cv2.imread(os.path.join(train_masks_path, x),\n", " cv2.IMREAD_ANYDEPTH)>0 for x in train_masks_filenames ]\n", "train_contours = [cv2.imread(os.path.join(train_contours_path, x),\n", " cv2.IMREAD_ANYDEPTH)>0 for x in train_contours_filenames ]\n", "\n", "# display the image\n", "plt.figure(figsize=(10,5))\n", "plt.subplot(1, 3, 1)\n", "plt.imshow( train_img[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Full-size training image' )\n", "# its \"mask\"\n", "plt.subplot(1, 3, 2)\n", "plt.imshow( train_masks[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Binary mask' )\n", "# and cell contours\n", "plt.subplot(1, 3, 3)\n", "plt.imshow( train_contours[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Object contour' )\n" ] }, { "cell_type": "markdown", "metadata": { "id": "WrZJHL0AkYzz" }, "source": [ "To facilitate their processing, we concatenate the binary masks and the contours to get one array with the training data." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "zv6cer-0klHJ" }, "outputs": [], "source": [ "import numpy as np\n", "train_output = [np.transpose(np.array([train_masks[i], train_contours[i]]),\n", " [1,2,0]) for i in range(len(train_masks))]" ] }, { "cell_type": "markdown", "metadata": { "id": "lPZVNfxDgubX" }, "source": [ "Next, we read images that will be used for validation and inspect some of them visually." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "VUkhuAHwgrqR" }, "outputs": [], "source": [ "# Path to the validation images\n", "val_input_path = 'dataset/validation_input'\n", "val_masks_path ='dataset/validation_binary_masks'\n", "val_contours_path = 'dataset/validation_contours'\n", "\n", "# Read the list of file names and sort them to have a match between images and masks\n", "val_input_filenames = [x for x in os.listdir(val_input_path ) if x.endswith(\".tif\")]\n", "val_input_filenames.sort()\n", "val_masks_filenames = [x for x in os.listdir(val_masks_path ) if x.endswith(\".tif\")]\n", "val_masks_filenames.sort()\n", "val_contours_filenames = [x for x in os.listdir(val_contours_path ) if x.endswith(\".png\")]\n", "val_contours_filenames.sort()\n", "\n", "print( 'Images loaded: ' + str( len(val_input_filenames)) )\n", "\n", "# Read training images\n", "val_img = [cv2.imread(os.path.join(val_input_path, x), cv2.IMREAD_ANYDEPTH) for x in val_input_filenames ]\n", "val_masks = [cv2.imread(os.path.join(val_masks_path, x), cv2.IMREAD_ANYDEPTH)>0 for x in val_masks_filenames ]\n", "val_contours = [cv2.imread(os.path.join(val_contours_path, x), cv2.IMREAD_ANYDEPTH)>0 for x in val_contours_filenames ]\n", "\n", "# Display the image\n", "plt.figure(figsize=(10,5))\n", "plt.subplot(1, 3, 1)\n", "plt.imshow(val_img[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Full-size training image' )\n", "# Its \"mask\"\n", "plt.subplot(1, 3, 2)\n", "plt.imshow(val_masks[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Binary mask' )\n", "# And cell contours\n", "plt.subplot(1, 3, 3)\n", "plt.imshow(val_contours[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Object contour' )" ] }, { "cell_type": "markdown", "metadata": { "id": "A7craIKDky78" }, "source": [ "Again, to facilitate their processing, we concatenate the binary masks and the contours to get one array with the training data." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "b1ea-0zskuNK" }, "outputs": [], "source": [ "# concatenate binary masks and contours\n", "val_output = [np.transpose(np.array([val_masks[i],val_contours[i]]),\n", " [1,2,0]) for i in range(len(val_masks))]" ] }, { "cell_type": "markdown", "metadata": { "id": "8R8Izj7YsLe3" }, "source": [ "## Preparing the training data\n", "Now, we are going to create the training set by randomly cropping the input images (and their corresponding labels) into **patches of 256 x 256 pixels**.\n", "\n", "All the images entering the network need to be normalized so their intensities are all of them in the range [0,1]. For this, we will create a function that normalizes the intensities values.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "WYmxsOS-875e" }, "outputs": [], "source": [ "import numpy as np\n", "\n", "# Function to create random patches\n", "def create_random_patches(imgs, masks, num_patches, shape):\n", " '''Create a list of image patches out of a list of images.\n", " Args:\n", " imgs (list): input images.\n", " masks (list): binary masks (output images) corresponding to imgs.\n", " num_patches (int): number of patches for each image.\n", " shape (2D array): size of the patches. Example: [256, 256].\n", "\n", " Returns:\n", " list of image patches and patches of corresponding labels (background, foreground and contours)\n", " '''\n", " input_patches = []\n", " output_patches = []\n", "\n", " height, width = shape\n", "\n", " for img, mask in zip(imgs, masks):\n", " img_height, img_width = img.shape[:2]\n", "\n", " for _ in range(num_patches):\n", " # Randomly select the top-left corner for cropping\n", " x = np.random.randint(0, img_width - width)\n", " y = np.random.randint(0, img_height - height)\n", "\n", " # Extract the patch from both the image and its corresponding mask\n", " input_patch = img[y:y+height, x:x+width]\n", " output_patch = mask[y:y+height, x:x+width]\n", "\n", " input_patches.append(input_patch)\n", " output_patches.append(output_patch)\n", "\n", " return input_patches, output_patches\n", "\n", "# Use the method to create six 256x256 pixel-sized patches per image\n", "train_input_patches, train_output_patches = create_random_patches(train_img, train_output, 6, [256, 256])\n", "\n", "# Normalization function using percentiles\n", "def normalizePercentile(x, pmin=1, pmax=99.8, axis=None, eps=1e-20):\n", " \"\"\"Percentile-based image normalization.\"\"\"\n", " mi = np.percentile(x, pmin, axis=axis, keepdims=True)\n", " ma = np.percentile(x, pmax, axis=axis, keepdims=True)\n", "\n", " # Normalize the image and ensure the values fall between 0 and 1\n", " x = (x - mi) / (ma - mi + eps)\n", " return np.clip(x, 0, 1) # Clip values to be between 0 and 1\n", "\n", "# Normalize the training input patches\n", "X_train = [normalizePercentile(x) for x in train_input_patches] # Normalize between 0 and 1\n" ] }, { "cell_type": "markdown", "metadata": { "id": "U2wIZabm8q6J" }, "source": [ "\n", "To follow Tensorflow standards, the input and output of the network have to be reshaped to 256 x 256 x 1. Therefore, the array containing the input images should have shape `[n, 256, 256, 1]` where `n` is the number of patches that will be used for the training.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "7zm1AADH9GmT" }, "outputs": [], "source": [ "# In X_train we will store the input images\n", "X_train = np.expand_dims(X_train, axis=-1)\n", "print('There are {} patches to train the network'.format(len(X_train)))" ] }, { "cell_type": "markdown", "metadata": { "id": "QXeRPKho9P4G" }, "source": [ "We want to predict three labels (background, foreground and cell-contours). One way to provide the data is with one hot encoding (e.g.,\n", "\n", "```\n", "[1, 2, 3, 2] --> [[1, 0, 0, 0],\n", " [0, 1, 0, 1],\n", " [0, 0, 1, 0]]\n", "```\n", "\n", "Create the ground truth array `Y_train` with a **one hot encoding format**. Note that you will need as many matrices as labels (i.e., 3) for each image." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "REisNUe69JJf" }, "outputs": [], "source": [ "def one_hot_encode(masks, num_classes=3):\n", " '''\n", " Convert label masks to one-hot encoding format.\n", "\n", " Args:\n", " masks (list or array): list of label masks (each mask should have pixel values corresponding to 3 classes).\n", " num_classes (int): number of classes (background, foreground, and contours).\n", "\n", " Returns:\n", " numpy array: One-hot encoded labels of shape (num_images, height, width, num_classes)\n", " '''\n", " # Check if masks are 2D or 3D and handle accordingly\n", " if len(masks[0].shape) == 2:\n", " height, width = masks[0].shape\n", " elif len(masks[0].shape) == 3:\n", " height, width, _ = masks[0].shape # Ignore the channel dimension if present\n", "\n", " # Create a placeholder for the one-hot encoded labels\n", " Y_train = np.zeros((len(masks), height, width, num_classes), dtype=np.uint8)\n", "\n", " for i, mask in enumerate(masks):\n", " # If the mask has 3 dimensions, we check if the number of channels is equal to the number of classes, otherwise print an error\n", " if mask.ndim == 3:\n", " if mask.shape[-1] == num_classes:\n", " Y_train[i] = mask\n", " else:\n", " # Iterate over each class\n", " Y_train[i, :, :, 1] = mask[:, :, 0] # Foreground (cells)\n", " Y_train[i, :, :, 2] = mask[:, :, 1]\n", " Y_train[i, :, :, 0] = 1 - (mask[:, :, 0] + mask[:, :, 1])\n", " else:\n", " # Assign each class to its corresponding channel\n", " Y_train[i, :, :, 0] = (mask == 0) # Background\n", " Y_train[i, :, :, 1] = (mask == 1) # Foreground (cells)\n", " Y_train[i, :, :, 2] = (mask == 2) # Contours\n", "\n", " return Y_train\n", "\n", "# Use the function to create the one-hot encoded ground truth\n", "Y_train = one_hot_encode(train_output_patches, num_classes=3)\n" ] }, { "cell_type": "markdown", "metadata": { "id": "GkWoS78u9L7_" }, "source": [ "Display the results" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "-dQj-_hDFEv4" }, "outputs": [], "source": [ "# Display one patch\n", "plt.figure(figsize=(25,5))\n", "plt.subplot(1, 5, 1)\n", "plt.imshow( X_train[0,:,:,0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Training patch' )\n", "# Background class\n", "plt.subplot(1, 5, 2)\n", "plt.imshow( Y_train[0,:,:,0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Binary patch (ground truth)' )\n", "# Foreground class\n", "plt.subplot(1, 5, 3)\n", "plt.imshow( Y_train[0,:,:,1], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Binary patch (ground truth)' )\n", "# Object contours\n", "plt.subplot(1, 5, 4)\n", "plt.imshow( Y_train[0,:,:,2], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Object contours patch (ground truth)' )\n", "# Reversed one hot representation\n", "plt.subplot(1, 5, 5)\n", "plt.imshow( np.argmax(Y_train[0], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Unique labelling of each pixel (ground truth)' )" ] }, { "cell_type": "markdown", "metadata": { "id": "YSDrFIzXVEF9" }, "source": [ "And now we do the same for the validation images:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "iBYhGigRVGbx" }, "outputs": [], "source": [ "# We first create the validation patches\n", "val_input_patches, val_output_patches = create_random_patches( val_img, val_output,\n", " 6, [256,256])\n", "\n", "# In X_val we will store the input images\n", "X_val = [normalizePercentile(x) for x in val_input_patches] # normalize between 0 and 1\n", "X_val = np.expand_dims(X_val, axis=-1)\n", "\n", "print('There are {} patches to validate the network'.format(len(X_val)))\n", "\n", "# In Y_val we will store the one-hot respresentation of the labels\n", "Y_val = [np.stack([1 - x[:,:,0] - x[:,:,1], x[:,:,0], x[:,:,1]],\n", " axis=-1) for x in val_output_patches ]\n", "Y_val = np.asarray(Y_val)\n", "\n", "# Display one patch\n", "plt.figure(figsize=(25,5))\n", "plt.subplot(1, 5, 1)\n", "plt.imshow( X_val[0,:,:,0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Validation patch' )\n", "# Background class\n", "plt.subplot(1, 5, 2)\n", "plt.imshow( Y_val[0,:,:,0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Binary patch (ground truth)' )\n", "# Foreground class\n", "plt.subplot(1, 5, 3)\n", "plt.imshow( Y_val[0,:,:,1], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Binary patch (ground truth)' )\n", "# Object contours\n", "plt.subplot(1, 5, 4)\n", "plt.imshow( Y_val[0,:,:,2], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Object contours patch (ground truth)' )\n", "# Reversed one hot representation\n", "plt.subplot(1, 5, 5)\n", "plt.imshow( np.argmax(Y_val[0], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Unique labelling of each pixel (ground truth)' )" ] }, { "cell_type": "markdown", "metadata": { "id": "H-M9o_URSzGJ" }, "source": [ "🌞 **Questions for discussion:**\n", "- In this exercise we are loading the entire dataset in memory. Do you think this is an optimal solution for a real application? Specially when we have more than 100 images.\n", "- Why do you think that we need to crop patches?\n", "- How do you decide the size of the pathces?\n", "- Explain in your own words what one-hot encoding is and why it is used in this notebook. How is it different from a simple binary representation of the labels?\n", "- Why do we need to normalize the intensity values of the input image? Can you think about would it be done for a different image modality (RGB) or for example, a whole slide images? Moreover, how can this affect the performance of the network with new images that have a different size or for which we did not crop patches?\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "erODPtqdS1UI" }, "source": [ "---\n", "⭐ Double click to write down your observation here\n", "\n", "\n", "```\n", "Answer:\n", "\n", "\n", "\n", "\n", "```\n", "---" ] }, { "cell_type": "markdown", "metadata": { "id": "BRNINwivV3km" }, "source": [ "## Custom segmentation metric: Jaccard index\n", "We define as well the [Jaccard index](https://en.wikipedia.org/wiki/Jaccard_index) (also known as Intersection over the Union or IoU) to monitor de segmentation performance.\n", "\n", "**Note**: by default we skip the background label in the calculation since most of the pixels are background and the metric value would be artificially high otherwise." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "CnDDhxpxoSkN" }, "outputs": [], "source": [ "import tensorflow as tf\n", "\n", "def jaccard_index( y_true, y_pred, skip_background=True ):\n", " ''' Define Jaccard index for multiple labels.\n", " Args:\n", " y_true (tensor): ground truth masks.\n", " y_pred (tensor): predicted masks.\n", " skip_background (bool, optional): skip 0-label from calculation.\n", " Return:\n", " jac (tensor): Jaccard index value\n", " '''\n", " # We read the number of classes from the last dimension of the true labels\n", " num_classes = tf.shape(y_true)[-1]\n", " # One_hot representation of predicted segmentation after argmax\n", " y_pred_ = tf.one_hot(tf.math.argmax(y_pred, axis=-1), num_classes)\n", " y_pred_ = tf.cast(y_pred_, dtype=tf.int32)\n", " # y_true is already one-hot encoded\n", " y_true_ = tf.cast(y_true, dtype=tf.int32)\n", " # Skip background pixels from the Jaccard index calculation\n", " if skip_background:\n", " y_true_ = y_true_[...,1:]\n", " y_pred_ = y_pred_[...,1:]\n", "\n", " TP = tf.math.count_nonzero(y_pred_ * y_true_)\n", " FP = tf.math.count_nonzero(y_pred_ * (y_true_ - 1))\n", " FN = tf.math.count_nonzero((y_pred_ - 1) * y_true_)\n", "\n", " jac = tf.cond(tf.greater((TP + FP + FN), 0), lambda: TP / (TP + FP + FN),\n", " lambda: tf.cast(0.000, dtype='float64'))\n", "\n", " return jac" ] }, { "cell_type": "markdown", "metadata": { "id": "let4jXdVtdQs" }, "source": [ "## Network definition\n", "Next, we define our U-Net-like network, with 3 resolution levels in the contracting path, a bottleneck, and 3 resolution levels in the expanding path.\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "K3O3hx36FxwF" }, "outputs": [], "source": [ "# Create U-Net for segmentation\n", "from tensorflow import keras\n", "from tensorflow.keras.models import Model\n", "from tensorflow.keras.layers import Input, UpSampling2D\n", "from tensorflow.keras.layers import Dropout\n", "from tensorflow.keras.layers import Conv2D, Conv2DTranspose\n", "from tensorflow.keras.layers import AveragePooling2D\n", "from tensorflow.keras.layers import concatenate\n", "from tensorflow.keras.optimizers import Adam\n", "# We leave the height and width of the input image as \"None\" so the network can\n", "# later be used on images of any size.\n", "inputs = Input((None, None, 1))\n", "\n", "# Contracting path\n", "c1 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (inputs)\n", "c1 = Dropout(0.1) (c1)\n", "c1 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c1)\n", "p1 = AveragePooling2D((2, 2)) (c1)\n", "\n", "c2 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p1)\n", "c2 = Dropout(0.2) (c2)\n", "c2 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c2)\n", "p2 = AveragePooling2D((2, 2)) (c2)\n", "\n", "c3 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p2)\n", "c3 = Dropout(0.3) (c3)\n", "c3 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c3)\n", "p3 = AveragePooling2D((2, 2)) (c3)\n", "\n", "# Bottleneck\n", "c4 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p3)\n", "c4 = Dropout(0.4) (c4)\n", "c4 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c4)\n", "\n", "# Expanding path\n", "u5 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same') (c4)\n", "u5 = concatenate([u5, c3])\n", "c5 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u5)\n", "c5 = Dropout(0.3) (c5)\n", "c5 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c5)\n", "\n", "u6 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same') (c5)\n", "u6 = concatenate([u6, c2])\n", "c6 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u6)\n", "c6 = Dropout(0.2) (c6)\n", "c6 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c6)\n", "\n", "u7 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same') (c6)\n", "u7 = concatenate([u7, c1], axis=3)\n", "c7 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u7)\n", "c7 = Dropout(0.1) (c7)\n", "c7 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c7)\n", "\n", "# The output will consist of 3 neurons (one per class) with softmax activation\n", "# so they represent probabilities\n", "outputs = Conv2D(3, (1, 1), activation='softmax') (c7)\n", "\n", "model = Model(inputs=[inputs], outputs=[outputs])\n", "\n", "model.summary()" ] }, { "cell_type": "markdown", "metadata": { "id": "fJkyiADzuro5" }, "source": [ "## Training the network\n", "Now we are almost ready to train our network! Some important training parameters to take into account:\n", "* `Epochs`: which defines the maximum number of epochs the model will be trained. Initially set to 100.\n", "* `Patience`: number of epochs that produced the monitored quantity (validation Jaccard index) with no improvement after which training will be stopped. Initially set to 50.\n", "* `Batch size`: the number of training examples in one forward/backward pass. Initially set to 10.\n", "* `Learning rate`: the parameter that determines the step size to update the weights of the network at each iteration while training the model.\n", "\n", "\n", "Since we have more than 2 output classes, we use the categorical cross-entropy (CCE) between the expected and the predicted pixel values as the loss function, and we also include the Jaccard index as a control metric. The Jaccard index obtained for the validation set during the training is taken into account to define an early stopping schedule for the training. This is important when you want to control the training time. Otherwise, you can wait until the network finishes the training and store the weights of the checkpoint that was doing better in the validation dataset.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "xycY-tS9F4PN" }, "outputs": [], "source": [ "from tensorflow.keras.callbacks import EarlyStopping\n", "\n", "# Training parameters\n", "numEpochs = 5 # suggested maximum number of epochs to train is 100\n", "patience = 5 # number of epochs to wait before stopping if no improvement\n", "batchSize = 10 # number of samples per batch\n", "lr = 0.0003 # learning rate\n", "# Define early stopper to finish the training when the network does not improve\n", "earlystopper = EarlyStopping(patience=patience, verbose=1, restore_best_weights=True,\n", " monitor='val_jaccard_index', mode='max')\n", "\n", "\n", "# Finally compile the model with Adam as optimizer, CCE as loss function and IoU\n", "# as metric\n", "opt = Adam(learning_rate=lr) # Adam with specified learning rate\n", "model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=[jaccard_index])\n", "\n", "# Train!\n", "history = model.fit( X_train, Y_train, validation_data = (X_val, Y_val),\n", " batch_size = batchSize, epochs=numEpochs,\n", " callbacks=[earlystopper])\n", "\n", "# Save the model weights to a HDF5 file\n", "model.save_weights('unet_pancreatic_cell_segmentation_best.weights.h5')\n", "\n", "# Save the model to a HDF5 file\n", "model.save('unet_pancreatic_cell_segmentation_model.h5')" ] }, { "cell_type": "markdown", "metadata": { "id": "i6d_arskwxuP" }, "source": [ "\n", "We can now plot the loss and metric curves for the training and validation sets.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "P_VsQDKUGSVS" }, "outputs": [], "source": [ "plt.figure(figsize=(14,5))\n", "\n", "# Summarize history for loss\n", "plt.subplot(1, 2, 1)\n", "plt.plot(history.history['loss'])\n", "plt.plot(history.history['val_loss'])\n", "plt.title('model loss')\n", "plt.ylabel('loss')\n", "plt.xlabel('epoch')\n", "plt.legend(['train', 'val'], loc='upper left')\n", "\n", "# Summarize history for Jaccard index\n", "plt.subplot(1, 2, 2)\n", "plt.plot(history.history['jaccard_index'])\n", "plt.plot(history.history['val_jaccard_index'])\n", "plt.title('model Jaccard index')\n", "plt.ylabel('Jaccard index')\n", "plt.xlabel('epoch')\n", "plt.legend(['train_jacc', 'val_jacc'], loc='lower right')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "id": "SYXO6kZfb0nO" }, "source": [ "🌞 **Exercise**: Play with the batch size and the learning rate. For example, check whether the network learns anything when the learning rate is too high (e.g., 0.01) or whether batch size has any effect on the final result. For this you can reduce the number of epochs so you do not need to wait for long training times.\n" ] }, { "cell_type": "markdown", "metadata": { "id": "HuzLtbRqPw_5" }, "source": [ "---\n", "⭐ Double click to write down your observation here\n", "\n", "\n", "```\n", "Answer:\n", "\n", "\n", "\n", "\n", "```\n", "---" ] }, { "cell_type": "markdown", "metadata": { "id": "SpMeDioSw406" }, "source": [ "## Check performance in the test set\n", "Finally we can load some test images for testing.\n", "\n", "🌞 **Exercises:**\n", "1. Load all the test images as it was done for the training and validation. Note that you will need to preprocess them as well. Call the variables `X_test` and `Y_test`. Also note that the original images have a size that is not divisible by a factor of 8 so the decoder side of the U-Net cannot create patches of the same shape as the input. Do you think that the network can process something that has a shape different from [256, 256]? how did we make this posssible in the architecture?\n", "\n", "2. On a paper, try computing the image shape resulting through each of the branches of the U-Net when down-sampling. You should see that if an image has input shape [234,234] for example, you will have some incompatibilities with the skip connections.\n", "Generally one would create a tiling strategy to reconstruct the entire images.\n", "To make it simpler, we will crop a patch as large as possible that is still divisible by 8.\n", "\n", "We can evaluate the network performance in test using both the CCE and Jaccar index.\n" ] }, { "cell_type": "markdown", "metadata": { "id": "LOPVk3e2xARB" }, "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "dvVKgUiRL4MJ" }, "outputs": [], "source": [ "# Evaluate trained network on test images\n", "\n", "# ⭐ Write your code here\n", "\n", "# Normalize input images\n", "X_test =\n", "# One-hot label representation\n", "Y_test =\n", "\n", "# Evaluate the model on the test data using `evaluate`\n", "results = model.evaluate(X_test, Y_test, batch_size=1)\n", "print('test loss CCE: {0}, Jaccard index: {1}'.format(results[0], results[1]))\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "c8vLUfOCxO9U" }, "source": [ "And also display some patches for qualitative evaluation." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "4o0P2LIGNqYn" }, "outputs": [], "source": [ "print('\\n# Generate predictions for 3 samples')\n", "predictions = model.predict(X_test[:1], batch_size=1)\n", "masks_prediction = np.array((predictions[0,:,:,0], predictions[0,:,:,1]))\n", "contours_prediction = predictions[0,:,:,2]\n", "print('predictions shape:', predictions.shape)\n", "\n", "# Display corresponding first 3 patches\n", "plt.figure(figsize=(48,15))\n", "plt.subplot(3, 7, 1)\n", "plt.imshow( X_test[0], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Test patch at low resolution' )\n", "# Side by side with its \"ground truth\"\n", "plt.subplot(3, 7, 2)\n", "plt.imshow( np.argmax(Y_test[0], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Ground truth' )\n", "# One hot final segmentation with argmax\n", "plt.subplot(3, 7, 3)\n", "plt.imshow(np.argmax(predictions[0], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Final segmentation segmentation' )\n", "# Foreground prediction\n", "plt.subplot(3, 7, 4)\n", "plt.imshow( predictions[0,:,:,1], 'CMRmap')\n", "plt.axis('off')\n", "plt.title( 'Foreground prediction' )\n", "# Contours predictions\n", "plt.subplot(3, 7, 5)\n", "plt.imshow( predictions[0,:,:,2], 'CMRmap')\n", "plt.axis('off')\n", "plt.title( 'Contours Prediction' )\n", "# ZOOM test ground truth\n", "plt.subplot(3, 7, 6)\n", "plt.imshow( np.argmax(Y_test[0,100:200, 100:200], axis=-1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Zoomed ground truth' )\n", "# ZOOM Ground Truth\n", "plt.subplot(3, 7, 7)\n", "plt.imshow( np.argmax(predictions[0][100:200,100:200], axis = -1), 'CMRmap', interpolation='nearest')\n", "plt.axis('off')\n", "plt.title( 'Zoomed prediction' )\n", "\n", "\n", "predictions = model.predict(np.expand_dims(X_test[50], axis = 0), batch_size=1)\n", "\n", "plt.subplot(3, 7, 8)\n", "plt.imshow( X_test[50], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Test patch at low resolution' )\n", "# Side by side with its \"ground truth\"\n", "plt.subplot(3, 7, 9)\n", "plt.imshow( np.argmax(Y_test[50], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Ground truth' )\n", "# and one hot final segmentation with argmax\n", "plt.subplot(3, 7, 10)\n", "plt.imshow(np.argmax(predictions[0], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Final segmentation segmentation' )\n", "# foreground prediction\n", "plt.subplot(3, 7, 11)\n", "plt.imshow( predictions[0,:,:,1], 'CMRmap')\n", "plt.axis('off')\n", "plt.title( 'Foreground prediction' )\n", "# contours predictions\n", "plt.subplot(3, 7, 12)\n", "plt.imshow( predictions[0,:,:,2], 'CMRmap')\n", "plt.axis('off')\n", "plt.title( 'Contours Prediction' )\n", "# ZOOM test ground truth\n", "plt.subplot(3, 7, 13)\n", "plt.imshow( np.argmax(Y_test[50,100:200, 100:200], axis=-1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Zoomed ground truth' )\n", "#ZOOM Ground Truth\n", "plt.subplot(3, 7, 14)\n", "plt.imshow( np.argmax(predictions[0][100:200,100:200], axis = -1), 'CMRmap', interpolation='nearest')\n", "plt.axis('off')\n", "plt.title( 'Zoomed prediction' )\n", "\n", "predictions = model.predict(np.expand_dims(X_test[80], axis = 0), batch_size=1)\n", "\n", "plt.subplot(3, 7, 15)\n", "plt.imshow( X_test[80], 'gray' )\n", "plt.axis('off')\n", "plt.title( 'Test patch at low resolution' )\n", "# Side by side with its \"ground truth\"\n", "plt.subplot(3, 7, 16)\n", "plt.imshow( np.argmax(Y_test[80], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Ground truth' )\n", "# and one hot final segmentation with argmax\n", "plt.subplot(3, 7, 17)\n", "plt.imshow(np.argmax(predictions[0], axis = -1), 'CMRmap', interpolation='nearest' )\n", "plt.axis('off')\n", "plt.title( 'Final segmentation segmentation' )\n", "# foreground prediction\n", "plt.subplot(3, 7, 18)\n", "plt.imshow( predictions[0,:,:,1], 'CMRmap')\n", "plt.axis('off')\n", "plt.title( 'Foreground prediction' )\n", "# contours predictions\n", "plt.subplot(3, 7, 19)\n", "plt.imshow( predictions[0,:,:,2], 'CMRmap')\n", "plt.axis('off')\n", "plt.title( 'Contours Prediction' )\n", "# ZOOM test ground truth\n", "plt.subplot(3, 7, 20)\n", "plt.imshow( np.argmax(Y_test[80,100:200, 100:200], axis=-1), 'CMRmap', interpolation='nearest')\n", "plt.axis('off')\n", "plt.title( 'Zoomed ground truth' )\n", "#ZOOM Ground Truth\n", "plt.subplot(3, 7, 21)\n", "plt.imshow( np.argmax(predictions[0][100:200,100:200], axis = -1), 'CMRmap', interpolation='nearest')\n", "plt.axis('off')\n", "plt.title( 'Zoomed prediction' )" ] }, { "cell_type": "markdown", "metadata": { "id": "HA0vKdPq7qSA" }, "source": [ "## Build a Web App to Deploy Your Trained U-Net\n", "\n", "You’ve trained a U-Net for image segmentation. In this lab, you’ll package it as a simple web app so colleagues without programming experience can try it in a browser: upload a TIFF → see the segmentation result.\n", "\n", "### Get prepared\n", "You can learn this by discussing with your favourate chat assistant, chatGPT, or the GitHub Copilot, or Gemini CLI. You can ask questions such as:\n", " - Write a prompt which explains what you are trying to do, e.g. you want to learn the basic of web app development in Python with the aim of creating your own image segmentation app based on U-Net.\n", " - what is a web application? How does it set up normally? \n", " - What is frontend and backend. \n", " - Get familiarize yourself with popular web app server framework such as `FastAPI`. \n", " - And then for the frontend, get yourself familiarize with the basic of HTML, CSS and Javascript.\n", "\n", "### Create Instructions for your AI Agent\n", "\n", "Once you get the basic idea on how web app works, proceed with creating your `GEMINI.md` file under the `Module2` folder. The goal is to prepare instructions for your gemini cli agent to help you build the web app -- this is a crucial context engineering step. Please be patient on this step, you need to sit down and write very detailed instructions to provide enough context for the AI to help you. \n", "\n", "There are many existing examples you can learn from, for example: https://github.com/dontriskit/awesome-ai-system-prompts you can find many real world examples, browe it through and you will see examples such as [lovable prompt](https://github.com/dontriskit/awesome-ai-system-prompts/blob/main/Loveable/Prompt.md). You don't need to understand everything, but reading the text part will get you some basic understanding on how production-level prompt looks like.\n", "\n", "\n", "Now follow the [core principles](https://github.com/dontriskit/awesome-ai-system-prompts?tab=readme-ov-file#the-foundation-core-principles-of-agentic-prompts) outlined in the read me file of this repo. try to include the following details in your `GEMINI.md` using markdown syntaxt:\n", " - Set a role for the AI agent\n", " - Describe the context, explain clearly what you have done, what you want to achieve\n", " - **Importantly:** include the full file path to the trained U-Net model file, tell the AI agent the framework is tensorflow, and the input image format. Also mention that you have a notebook file which produced the model.\n", " - Tell it to use `fastapi` for the backend server\n", " - And for the frontend, you want to create a single page web app using CSS framework `tailwind CSS`. \n", " - Use AI assistant to polish your `GEMINI.md` file.\n", "\n", "\n", "### Start building\n", "\n", "Now move on with building the web app by starting the `gemini` CLI in your VS Code terminal.\n", "\n", "Tell it to create the web app to serve the U-Net model, and you want to be able to launch the server later, and be able to upload an tiff image and see the segmentation in your web browser.\n", "\n", "After building the app, you need to ask gemini the command to start the server. And you might need **Port Forwarding** (see below) to actually see the app in your browser.\n", "\n", "To test the app, here is [an example image](https://raw.githubusercontent.com/aicell-lab/ddls-course/main/static/uploads/ddls_2025_lab2_example_image.tif) you can download, then select in your U-Net app to show whether it works.\n", "\n", "If it doesn't work, tell what is the issue and what you want to change. Sometimes, you might need to use the browser developer tool to see error from the browser console to debug frontend code, e.g. [in Chrome](https://developer.chrome.com/docs/devtools/open). \n", "\n", "## Note on Port Forwarding\n", "\n", "Since you’re running your FastAPI server **on a remote machine** (e.g., a university server or a cloud VM), it’s not automatically visible on the public internet. By default, only you can reach it from inside the machine. To make your web app usable from your browser — and shareable with others — you need to create a **tunnel**. This process is called **port forwarding**.\n", "\n", "### What is port forwarding?\n", "\n", "* Your FastAPI server runs on a specific port (e.g., `http://localhost:8000`).\n", "* But “localhost” on the remote server is not the same as “localhost” on your laptop.\n", "* Port forwarding creates a secure bridge: traffic to a remote port is forwarded through VS Code to a special URL that you (and optionally others) can access in a browser.\n", "\n", "In short: **without port forwarding, your app stays hidden; with it, you get a shareable link.**\n", "\n", "### How to set it up in VS Code\n", "\n", "1. Start your FastAPI server inside the VS Code terminal:\n", "\n", " ```bash\n", " uvicorn backend.app:app --reload --port 8000\n", " ```\n", "2. In VS Code, look at the bottom panel tabs (“PROBLEMS”, “OUTPUT”, “TERMINAL”, etc.).\n", "3. Click on the **PORTS** tab.\n", "4. You should see port `8000` listed (if not, click ➕ and add it manually).\n", "5. In the PORTS panel, you’ll see a globe icon 🌐. Click it to open the forwarded URL in your browser.\n", "6. By default, the link is **Private** (only you). To share:\n", "\n", " * Right-click the port entry → **Change Port Visibility** → set to **Public**.\n", " * Now you’ll have a globally accessible URL to share with colleagues.\n", "\n", "⚠️ **Important:** Be mindful that making the app public means anyone with the link can use it. Do this only for demo purposes." ] } ], "metadata": { "accelerator": "GPU", "colab": { "provenance": [ { "file_id": "1kKVJZCEDjAU_FKet2re4jLieWnw8zgaY", "timestamp": 1756801245538 } ] }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.7.7" } }, "nbformat": 4, "nbformat_minor": 0 }