Image segmentation with a U-Net-like architecture - Keras Code Examples

Key Takeaways

welcome to the henry ai labs walkthrough of keras code examples keras has provided 56 code examples implementing popular ideas in deep learning this ranges from the basics such as simple mnist and imdb text classification all the way to cutting edge research ideas such as knowledge distillation supervised contrastive learning and transformers we'll also explore fun generative examples like variational autoencoders and cyclegan my contribution to these code examples is to explain every single line of code in each of them walking through each of the individual keras examples i'm not the author of these code examples please consider starting the github repositories to show support to the original authors in this video you'll learn how the image segmentation problem is set up to have an input image and then an output image in this unit like architecture that compresses the dimensionality of an image from say the original 160 by 160 feature planes down into 10 by 10 spatial resolution and then back up into 160 by 160. this is where it gets the u-net part of its name we'll also see how to run inference for given individual images for the semantic segmentation problem in the end we'll produce this pixel map for this image where this is our ground truth label y now that we've gotten the basic image classifier and mnist examples out of the way we'll graduate into more interesting and exciting applications of deep learning and computer vision this problem is image segmentation with a u-net-like architecture image segmentation is probably best known for self-driving cars where we want the cars to be able to label every pixel they're seeing at every time on the road this involves labeling the road pedestrians stop signs traffic lights all sorts of things that help the autonomous vehicle system have a sense of everything it's currently seeing so first we download the data set we're going to have these pet images and then we're going to have their annotations where every pixel has been labeled as either a part of the cat or whatever pet or a part of the background so we're doing a binary classification with each pixel in our image segmentation task to segment each pixel into being a member of the cat or a member of the background so note how here differently from the first example where we had a zip file here we have a tar.gz file so the only difference is instead of shell command unzip we do shell command tar dash xf and then images.tar.gz and then we have the data set in our working directory from this web link this next block of code is aligning the images with their annotations so one important thing about this is we need to insert this sorted line in order to align the original images with their corresponding annotations this is because something with python or some underlying thing that i'm not exactly sure of but it'll scramble these uh these numbers so it'll go 110 100 rather than one two three four five six seven eight nine so we need to use this sorted syntax so that it'll sort it alphabetically and align the correct original pet image with its annotation so all this code is doing is it's looping through the directory and then it is uh doing the path join and then sorting it such that it's uh numerically ascending so that you can easily align these two with each other so then what this is doing is zipping together the first 10 input images the first 10 target images and then printing out these two different paths as we traverse through our directory shown here on the left to get to these uh different instances in our data set so we go to images and then we have this folder it'll take a second to load because it's a really large folder with a lot of files in it and then see how we navigate from images slash abyssian underscore 1.jpg so this is how we're going to be looping through this to get these images and then align them and then pipeline them into our model this line of code will visualize an original cat image with its semantic segmentation ground truth label map so we use the pil dot image operations auto contrast in order to set the values to be either completely black or completely white when we're visualizing this pixel map so we see the original cat here we have labels for the border of the cat as well as the cat itself and then the background of the cat so we're going to be using a convolutional network that takes in this original image and outputs this image and this is the ground truth label y as we do supervised learning y minus y prime and apply a loss to each of these individual pixel predictions the next step is to prepare a sequence class to load and vectorize batches of data so the keras sequence class that we're overloading is going to be the way that we're loading these images into memory using the load image function and then pre-processing them and returning x y pairs to fit our model in this case the x are going to be images like this and the y's are going to be semantic segmentation mass like this so the sequence class as shown in the documentation here from tensorflow is one way of pipelining this data so this note describes probably the motivating reason for using sequences it says that they're a safer way to do multi-processing and the structure guarantees that the network will only train once on each sample per epoch which is not the case with say generators and i'm assuming the previous example of image dataset from directory was probably an example of a generator since they're making this distinction here although i could be wrong about that but what's important for our sake is noting that this is just a way of pipelining batches of data so if you're familiar with pytorch it's the same kind of syntax of initializing the dataset class having some function to get the length of data set and then having that this is probably the meat of it how you get the next item so in this case we originally defined the batch size the image size the input paths the target paths these are the paths as we traverse through this directory again we have this images folder and we also have the annotations folder that contains the pixel maps and this contains the original pet images and here what we're doing is we're defining a way to index the data set so some of the most important things about this are looping through the paths originally having this placeholder of mp.0s numpy.0s is a way to define an empty tensor so it's originally an empty tensor and then we overload this empty tensor by indexing it with our image so x subscript j is this image so we have this original x which is a placeholder for our entire batch so again we're constructing a tensor that's the size of the batch by the size of the image sizes and then we're loading in each of these images to fill out this batch and then when you do get item you get a batch of these pet images and their annotations so this does not just return one cat and one pixel map for the cat unless you had the batch size equal to be one otherwise it's trying to get a whole batch of cats and a whole batch of pixel maps for the cats or dogs or whatever's in this oxford pets data set so probably the most important thing to note is how they use these placeholders how they loop through the batches how they overwrite indexes in the original zeroed out tensor with these new values and then how they return x y pairs in this sequence data loader keras class so now that we have a data loader that returns batches of images and their corresponding pixel maps we define the unet exception style model so the exception style is defining is describing a particular way of structuring these feature planes and these skip connections but for now we'll just abstract this away as one type of convolutional neural network design so probably more interesting is the u-net design so these blocks of code again def define git model this will return a model with the given input image size and the number of classes steps through these separate layers of the keras layer construction passing in the previous output you see inputs goes to inputs x goes to x x goes to x same kind of thing all the way down same idea previously looping through this array of filter maps to step through and then construct this block loop again construct this block except passing this in where you have this filters parameter and then you construct this network and then again you have another loop that constructs this u-net like architecture so in model.summary we see this u-net architecture so what this is describing is how the spatial resolution is going to get down sampled from 160 by 160 all the way down to 10 by 10 and then back up to the output of 160 by 160 which is the corresponding output space of these pixel maps so it takes in 160 by 160 rgb say cat image and then it passes it through all these convolutional layers you know with the skip connections all this kind of underlying complexity with the convolutional network respective parameter accounts for each layer and then it starts up sampling it back into the target output size of these semantic pixel maps this code example is a way of manually constructing a validation split it's common practice in machine learning to have a train and a test set and additionally have a validation set so you usually have train validation and test and this is done to do hyper parameter tuning with the training set without corrupting the integrity of this by looking into the test set and then optimizing the hyper parameters to overfit to this test set so this is an example of how to have these image paths and then explicitly index them with this list syntax of colon then minus validation samples one way of having 1000 samples explicitly constructed from this data set to be used in our validation set and then again we use another example of this data sequence this is when we're actually initializing it with the train set so we've defined our sequence before up top how we define what it this class is our overwritten sequence to customize it for loading this particular kind of data and then this code below is an example of actually instantiating this and having the train generator that's an instance of this sequence data class then we also have our validation generator in this case we're inputting these different image paths so particularly what we do in this line of code is we say these are the image paths for the training set these are the image paths for the validation set and again the image paths is referencing how you traverse this directory because we've loaded these directories onto our disk with our original call to the curl shell command to access this data set from the internet and download it onto the disk in our collab directory this is where we have our original images and their ground truth annotation labels so now we're training the model we compile it with the rms prop optimizer algorithm we have the sparse categorical cross entropy loss on each pixel in the pixel map then we have the callback where we save these weights at each checkpoint that means each epoch we're going to look to save the weights but only if it's the best it's done so far this is an extra argument you pass into the keras model checkpoint callback so now we use model.fit and then we step through this data so in the same theme this is the speed running on the cpu and now we'll be right back re-running all this code on the gpu which will be much faster so now we've rerun all the code but except for this time it's on the gpu accelerator so we're loading this model onto the gpu this time instead of the cpu and we're noting how much faster it runs through each epoch of the data compared to running on the cpu so previously we're looking at 40 minutes to step through the data here we're looking at about 30 seconds so note how much quicker the gpu is running this deep learning training compared to running on a cpu 15 epochs of stepping our u-net exception-like model through the pairs of pet images and their corresponding pixel labels takes roughly about 15 minutes and this is the uh progress in the training we see the loss steadily decreasing on our validation set as well as our training set so in this case we have the training set the loss goes from two all the way down to 0.18 pretty much monotonically decreasing although our validation loss doesn't continue to decrease with our training loss so this is something we'll probably talk about more later data augmentation regularization all sorts of different stuff to try to make the validation loss decrease monotonically alongside the trading loss but see still we see this epoch progress report how long it takes per step each epoch and then the change in our training loss and our validation set loss and remembering again that we had pre that we had explicitly subsetted our data to make this validation set then we passed these image paths into our overwritten sequence data loading class so now that we have a trained model we can visualize some predictions from this model using this code so this function is displaying a mask so what it has to do is at the output of our neural network we have a soft max activation a softmax activation is a probability distribution so we have to take the arg max meaning the slot that has the most probability density to output the samples from this model when it has a soft max activation at the target this is different from things like say having a sigmoid at the target where you wouldn't need this kind of arg max so now we expand the dimensions of our mass so that it can be inputted to our pil image loading thing and then we pass in the mask the image array image thing and then we have the display and then we have the auto contrast that helps us visualize the corresponding pixel map so this up top is the ground truth label note how it does a much better job of highlighting the color of the dog compared to this which is the prediction from our model but we see that our model does do a really good job of outlining the dog it just isn't quite as precise as the label we also see it misses this back leg but this is the code that you need to visualize the semantic segmentation predictions from this model so to recap this is the first example in this keras code example series that looks through a really interesting problem of image segmentation we saw how we have the images and their corresponding pixel map annotations how we can visualize these pixel maps with the semantic labels how we can overwrite the sequence data loader class from keras in order to batch our data and pipeline x y batches into our model the x being these batches of pet images and the y's being their corresponding pixel maps we saw how the u-net architecture has this way of going from 160 by 160 down to 10 by 10 feature plane sizes and then back into the original dimensionality of the image to correspond a pixel labeling for each pixel in the original map which is done by having the same output size as the original input image as well as the corresponding annotation mask also has a size 160 by 160. then we saw how to explicitly index the path the list that contains all the paths by doing this indexing syntax and then how to pass this as an input to our overwritten sequence data loader we saw the dramatic speed up from 40 minutes down to 40 seconds when switching from the cpu to the gpu in order to train the semantic segmentation model another example of using the checkpoint to save the weights and then we saw an example of how to sample a prediction from this model on a given individual image and we saw how it does a pretty good job of getting this correct pixel map from this dog image [Music]