How To Run UNET On UbiOps

In this article:

1. Biomedical Image Segmentation
2. Why use AI for biomedical image segmentation?
3. What is UNET?
4. Why use UNET?
5. Run UNET on UbiOps for biomedical image segmentation
6. Results of the implementation

Artificial Intelligence (AI) has been revolutionizing the field of medicine in the past years. In 2022, the global market size of artificial intelligence in healthcare was valued at 15.4 billion USD by Grand View Research. Because AI can assist doctors and researchers with all kinds of relevant issues, this market is still expected to grow rapidly. Some examples of how AI is being applied in healthcare are:

• Analysis of medical images
• Drug discovery
• Personalized medicine
• Predictive analytics
• Robotic surgery

The integration of machine learning (ML) systems in the medical industry allows for accurate early prediction of diseases, and reliability in treatments in cases like surgery. An application that we will focus on in this article is the application of AI on medical images. It’s commonly used in the field of cancer detection. AI can analyze images from X-rays, CT scans, and MRIs, etc., and help physicians and researchers to identify and diagnose diseases more accurately and efficiently. More specifically, this article will focus on biomedical image segmentation.

Let’s Explore Biomedical Image Segmentation

The goal of biomedical image segmentation is to divide a medical image into different sections, where each section corresponds to a specific anatomical structure, or tissue type. These regions can then be quantified and analyzed, which helps in diagnosis and research. For example, say you have an X-ray image of a patient’s lungs, you might want to segment the lungs for analysis. It could look something like this:

Image segmentation can also be applied on microscope images. In fundamental cancer research images of cancer cells can be analyzed after the cells have been segmented. Overall, biomedical image segmentation is an important process that enables accurate and precise analysis of medical images, leading to improved diagnosis and research in a variety of fields. Binary masks like in these two examples are very convenient for automated analysis. However, you can imagine they can be very tedious to make by hand. In more complex cases, they require medical expertise and many hours of manual work to create.

Why use AI for biomedical image segmentation?

That is where AI comes in. A well trained algorithm can perform biomedical segmentation as accurately as a human, but much faster. With help from AI, an expert would only have to make the training data, and then the AI can do the rest. This means that AI can aid in processing large datasets that no human could ever process in a reasonable amount of time. A well trained AI is also less prone to error. Automated systems will be more consistent than a human. These systems can also be combined with others, like noise reduction, to create an even more accurate segmentation pipeline.

What is UNET?

One of the most popular image segmentation algorithms is called UNET, an algorithm in the class of convolutional neural networks that gets its name from the shape of its model architecture; UNET is short for U-shaped network. A UNET performs something called semantic segmentation. That means that it classifies every pixel from an input image. That is exactly what we need to perform biomedical image segmentation! For each pixel in an input image, we want to determine whether it is part of an object of interest or not.

So what about UNET is U-shaped, exactly? Its architecture. But what does that mean? The UNET architecture consists of two parts: an encoder and a decoder. The encoder is a series of convolutional layers that gradually reduce the spatial resolution of the input image while increasing the number of feature maps. The decoder, on the other hand, is a series of convolutional layers that gradually increase the spatial resolution of the feature maps while reducing the number of feature maps. The encoder and decoder of UNET allow it to capture both low-level and high-level features of an image, making it capable of distinguishing different regions of an image. When visualized, the U-shape of the architecture appears. The following schematic of the UNET architecture is from the original paper on UNET.

Why would I want to use UNET?

There are many good reasons to want to implement a UNET for your image segmentation tasks. Firstly, the UNET architecture is designed for image segmentation and it generally achieves better results than traditional non-machine learning based image analysis methods like the seeded watershed method. Its architecture is also quite simple and easy to implement, and examples of implementations can easily be found online. Therefore, you will have a working model up and running so you can start your analyses more quickly. Lastly, UNET is also computationally efficient, as it needs relatively few images to be sufficiently trained.

Either way, training and serving convolutional neural networks like UNET takes a lot of computation power. The cost of building and maintaining the computing infrastructure to perform these tasks is high, because GPUs and ML engineers are expensive! UbiOps can save you effort and money with our on-demand, scale-to-zero platform. Simply deploy your model on the platform, and our GPUs will do the computing for you. Deployment is easy, and we offer support in the deployment and maintenance of your model. The UbiOps platform is also very reliable and ISO certified, which is extremely important in healthcare. We have a large number of GPUs, so we are always ready to run your AI service.

UNET on UbiOps for biomedical image segmentation

So we can conclude UNET is an important machine learning architecture in medical technology, and UbiOps can help you run it reliably and cheaply. So let’s deploy one! To show how to run UNET on UbiOps, we will build a model for a simple cell segmentation case and deploy it to UbiOps.

What dataset will we apply it to?

We will use a sample of around 100 images randomly selected from this dataset from the Cell Image Library. These are images of U2OS cells, a cell line widely used in genomic instability research. To give you an idea of what these images look like: the microscope image above is from this dataset.

To train a neural network, you need ground truth. These are examples of perfect output for a given input, and have to be made manually by an expert. The image to the right of the microscope image is the ground truth for that image. In this specific case, the creation of ground truth is very simple. Because cell segmentation is a common problem, many tools exist to do just that. We used NucleAIzer to segment the cells, and then applied simple thresholding in ImageJ to create the binary masks that will be used to train UNET.

We realize the existence of these tools makes this UNET case trivial. The purpose of this article is to show how to use UbiOps to run AI for biomedical image processing or other UNET applications. The same approach will work for more complex cases, as well.

UNET implementation

The UNET implementation used here is a modified version of the implementation supplied in this article by Harshall Lamba, which also offers a more in-depth explanation of UNET if you’re interested. We changed the activation function: instead of ReLU, we opted for Leaky ReLU as it gave better results. Also, UNET has been shown to work just fine on 128×128 pixel images, so to save some time I wrote image cutting and stitching functions to increase the number of images without having to make more ground truth. I used the following two ImageJ macros to get my images to the right size, and create the ground truth masks from the NuclAI output.

setBatchMode(true);
input = "C:/<FILE LOCATION>/";
list = getFileList(input);

for (i = 0; i < list.length; i++) {
	convert(input, list[i]);
}
setBatchMode(false);

function convert(input, filename) {
	open(input + filename);
	makeRectangle(0, 0, 512, 512);
	run("Crop");
	run("8-bit");
	run("Save");
	close();
}

setBatchMode(true);
input = "C:/<FILE LOCATION>/";
list = getFileList(input);

for (i = 0; i < list.length; i++) {
	convert(input, list[i]);
}
setBatchMode(false);

function convert(input, filename) {
	open(input + filename);
	run("8-bit");
	setAutoThreshold("Default dark no-reset");
	//run("Threshold...");
	//setThreshold(45, 255);
	setOption("BlackBackground", true);
	run("Convert to Mask");
	run("Fill Holes");
	saveAs("Tiff", input + filename + ".tif");
	close();
}

After all preprocessing, we ended up with a dataset of 1264 128×128 pixel images, that were partitioned into a training set and validation set at a 7:1 ratio. The model was trained for 52 epochs with a batch size of 16. It’s possible to train a model on UbiOps. In this case I trained the model locally and put the trained model in my deployment on UbiOps.

Deploying UNET to UbiOps

So now we have our model and want to deploy it to UbiOps for inference. UbiOps creates a scalable inference API endpoint. This means that once our model is deployed, we can use our UNET from anywhere and integrate it with the rest of our stack using the UbiOps API.

For UNET to work on UbiOps, we need to deploy our model, any supporting functions we use, and the code used to make predictions with the model. The code of the model architecture and the supporting functions we use in this case are found in the snippet below this paragraph. This model is borrowed, and slightly altered, from this Towards Data Science article by Harshall Lamba. The functions tile() and stitch() are added, and used to cut the microscope images into the right size for the model, and then stitch them back together after prediction. This way we can make predictions for any 128x128y pixel sized images.

Deploy model on UbiOps

import sys
import os
import numpy as np

from skimage.transform import resize
from itertools import product
from PIL import Image

import tensorflow as tf

from keras.models import Model
from keras.layers import Input, BatchNormalization, Dropout, LeakyReLU
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras.layers.pooling import MaxPooling2D
from keras.layers import concatenate
from tensorflow.keras.optimizers import Adam

def tile(filename, dir_in="", d=128):
    image_list = np.zeros((16, 128, 128, 1), dtype=np.float16)
    counter = 0
    name, ext = os.path.splitext(filename)
    img = Image.open(os.path.join(dir_in, filename))
    w, h = img.size
    
    grid = product(range(0, h-h%d, d), range(0, w-w%d, d))
    for i, j in grid:
        box = (j, i, j+d, i+d)
        cropped_im = img.crop(box)
        image = np.array(cropped_im, dtype=np.float32)
        image = resize(image, (128, 128, 1), mode = 'constant', preserve_range = True)
        image_list[counter] = image/255.0
        counter += 1
    
    return image_list


def stitch(image_list, d=512):
    ### stitches 128x128 images from dir_in into a square of size d
    blank_image = Image.new("L", (d, d))
    grid_size = int(d/128)
    counter = 0
    
    for i in range(grid_size):
        for j in range(grid_size):
            im = Image.fromarray(np.squeeze((image_list[counter]*255).astype(np.uint8), axis=2))
            blank_image.paste(im, (j*128, i*128))
            counter += 1
            
    return blank_image


def conv2d_block(input_tensor, n_filters, kernel_size = 3, batchnorm = True):
    """Function to add 2 convolutional layers with the parameters passed to it"""
    # first layer
    x = Conv2D(filters = n_filters, kernel_size = (kernel_size, kernel_size),
              kernel_initializer = 'he_normal', padding = 'same', activation=LeakyReLU(alpha=0.2))(input_tensor)
    if batchnorm:
        x = BatchNormalization()(x)
    
    # second layer
    x = Conv2D(filters = n_filters, kernel_size = (kernel_size, kernel_size),
              kernel_initializer = 'he_normal', padding = 'same', activation=LeakyReLU(alpha=0.2))(input_tensor)
    if batchnorm:
        x = BatchNormalization()(x)
    
    return x

    
def get_unet(input_img, n_filters = 16, dropout = 0.1, batchnorm = True):
    """Function to define the UNET Model"""
    # Contracting Path
    c1 = conv2d_block(input_img, n_filters * 1, kernel_size = 3, batchnorm = batchnorm)
    p1 = MaxPooling2D((2, 2))(c1)
    p1 = Dropout(dropout)(p1)
    
    c2 = conv2d_block(p1, n_filters * 2, kernel_size = 3, batchnorm = batchnorm)
    p2 = MaxPooling2D((2, 2))(c2)
    p2 = Dropout(dropout)(p2)
    
    c3 = conv2d_block(p2, n_filters * 4, kernel_size = 3, batchnorm = batchnorm)
    p3 = MaxPooling2D((2, 2))(c3)
    p3 = Dropout(dropout)(p3)
    
    c4 = conv2d_block(p3, n_filters * 8, kernel_size = 3, batchnorm = batchnorm)
    p4 = MaxPooling2D((2, 2))(c4)
    p4 = Dropout(dropout)(p4)
    
    c5 = conv2d_block(p4, n_filters = n_filters * 16, kernel_size = 3, batchnorm = batchnorm)
    
    # Expansive Path
    u6 = Conv2DTranspose(n_filters * 8, (3, 3), strides = (2, 2), padding = 'same')(c5)
    u6 = concatenate([u6, c4])
    u6 = Dropout(dropout)(u6)
    c6 = conv2d_block(u6, n_filters * 8, kernel_size = 3, batchnorm = batchnorm)
    
    u7 = Conv2DTranspose(n_filters * 4, (3, 3), strides = (2, 2), padding = 'same')(c6)
    u7 = concatenate([u7, c3])
    u7 = Dropout(dropout)(u7)
    c7 = conv2d_block(u7, n_filters * 4, kernel_size = 3, batchnorm = batchnorm)
    
    u8 = Conv2DTranspose(n_filters * 2, (3, 3), strides = (2, 2), padding = 'same')(c7)
    u8 = concatenate([u8, c2])
    u8 = Dropout(dropout)(u8)
    c8 = conv2d_block(u8, n_filters * 2, kernel_size = 3, batchnorm = batchnorm)
    
    u9 = Conv2DTranspose(n_filters * 1, (3, 3), strides = (2, 2), padding = 'same')(c8)
    u9 = concatenate([u9, c1])
    u9 = Dropout(dropout)(u9)
    c9 = conv2d_block(u9, n_filters * 1, kernel_size = 3, batchnorm = batchnorm)
    
    outputs = Conv2D(1, (1, 1), activation='sigmoid')(c9)
    model = Model(inputs=[input_img], outputs=[outputs])
    return model

We also want to deploy code to actually use our model. This is done using a Deployment class that UbiOps uses to handle model initialization and requests. The code is in the snippet below. The __init__() function runs when your model is initialized on UbiOps, and the request() function runs when you create a request to use your model. The template for this code is found in the documentation. You only need to add the code that is unique to your use case. In ours, the initialization function loads our model and the model weights that we found in training. The request function uses Tensorflow to create a prediction. And that’s it!

class Deployment:

    def __init__(self, base_directory, context):
        """
        Initialisation method for the deployment. It can for example be used for loading modules that have to be kept in
        memory or setting up connections. Load your external model files (such as pickles or .h5 files) here.

        """

        print("Initialising the model")

        model_file = os.path.join(base_directory, "model.h5")
        input_img = Input((128, 128, 1), name='img')
        self.model = get_unet(input_img, n_filters=16, dropout=0.05, batchnorm=True)
        self.model.compile(optimizer=Adam(learning_rate=0.005), loss="binary_crossentropy", metrics=["accuracy"])
        self.model.load_weights('model.h5')


    def request(self, data):
        """
        Method for deployment requests, called separately for each individual request.

        """
        print('Loading data')
        image_list = tile(data['image'])
        
        print("Prediction being made")
        pred = self.model.predict(image_list, verbose=1, batch_size = 32)
        pred_t = (pred > 0.50).astype(np.uint8)
        prediction = stitch(pred_t)
        
        rel_output_path = "prediction.jpg"
        prediction.save(rel_output_path, "JPEG", quality=100, optimize=True, progressive=True)
                
        return {
            "jpg": rel_output_path,
        }

Create local folder

Now, to get all of this onto UbiOps, we will create a local folder that contains deployment.py, our file of model weights, and a requirements.txt. The Python script contains the code we want to deploy. The requirements.txt contains all the libraries that UbiOps needs to install for your code to work. To create this local folder we will use a tutorial that is found in the documentation. You can find the code for this specific case below.

API_TOKEN = 'Token MYTOKEN' # Make sure this is in the format "Token token-code"
PROJECT_NAME = 'unet-op-ubiops' #CHANGE THIS TO YOUR OWN PROJECT NAME
DEPLOYMENT_NAME = 'unet'
DEPLOYMENT_VERSION = '1'

# Import all necessary libraries
import shutil
import os
import ubiops

client = ubiops.ApiClient(ubiops.Configuration(api_key={'Authorization': API_TOKEN}, 
                                               host='https://api.ubiops.com/v2.1'))
api = ubiops.CoreApi(client)

# Initiate a local directory
os.mkdir('UNET_deployment_package')


%%writefile UNET_deployment_package/requirements.txt

tensorflow==2.10.0
numpy==1.23.5
tifffile==2023.2.28
scikit-image==0.18.3
scikit-learn==1.1.1
natsort==7.1.1
Pillow==8.4.0
Keras==2.10.0
keras-preprocessing==1.1.2


# Create the deployment
deployment_template = ubiops.DeploymentCreate(
    name="unet",
    description='UNET deployment',
    input_type='structured',
    output_type='structured',
    input_fields=,
    output_fields=,
    labels={"demo": "UNET"}
)

api.deployments_create(
    project_name="unet-op-ubiops",
    data=deployment_template
)

# Create the version
version_template = ubiops.DeploymentVersionCreate(
    version="v1",
    language='python3.9',
    instance_type='512mb',
    minimum_instances=0,
    maximum_instances=1,
    maximum_idle_time=1800, # = 30 minutes
    request_retention_mode='none' # we don't need request storage in this example
)

api.deployment_versions_create(
    project_name="unet-op-ubiops",
    deployment_name="unet",
    data=version_template
)

# Zip the deployment package
shutil.make_archive('UNET_deployment_package', 'zip', '.', 'UNET_deployment_package')

# Upload the zipped deployment package
file_upload_result =api.revisions_file_upload(
    project_name="unet-op-ubiops",
    deployment_name="unet",
    version="v1",
    file='UNET_deployment_package.zip'
)

client.close()

When we run this code in a Jupyter notebook that also contains our functions and Deployment class, everything will be automatically deployed to UbiOps. It will be ready for use in only a few minutes!

Building your deployment from the provided code can take a couple minutes. To save time, you can test your deployment locally. This will give you the exact errors that UbiOps will give you when you try to deploy, but it’s much faster! I use this to iron out some mistakes.

After fixing a couple bugs, our UNET has been deployed and the build is successful! Now we can create a request to use our UNET model. In this build, we use UbiOps’s storage feature, but it’s also possible to use other cloud storage like Google Cloud Storage. To create our request, we will go to our deployment on UbiOps and create it manually. It’s also possible to create a request from a Python client!

We navigate to our deployment on the UbiOps platform, and press Create Request. We will then be prompted to select an input file, and UbiOps will run our model and store the output in a location that we have specified in the deployment.

Results

Now, let’s create a request in UbiOps with a completely new image that the AI has never seen before. You can see it works very well!

Conclusion

UNET is a commonly used neural network architecture in biomedical image segmentation. We hope this article made it clear how you can use UbiOps to quickly and reliably run this machine learning model in the cloud.

If you have any questions about what else is possible with UbiOps, or you’re interested in deploying your model to UbiOps, do not hesitate to reach us and book a call with our Product Specialist!