Learn how to remove image backgrounds with PyTorch - Deep Learning Tutorial

hi everyone I'm Patrick from assembly Ai and in this tutorial we learn how we can remove the background of images with pytorch so here we have an input image and then after applying the model we can extract only the foreground and then we can also merge this with another background that we want so without further Ado let's jump right into it so before getting started I want to give credits to this repository that I found on GitHub it's called practical ML and contains a lot of very cool deep learning tutorials including this one about removing the image background so a lot of the code is taken from this repository I just did some slight modifications so you have full credits to this author and then the model we are going to use it can be found on the pi torch up and it's the Deep lab version 3 image segmentation model so we're not gonna develop the model from scratch we just used this model from The Hub and yeah so let's get started so we're going to implement this in a Google collab I put the link in the description below and then you can grab this and follow me here so the first thing you want to do is change the runtime type to a GPU and then you want to upload your file your image and then here you change the input file name and then you can also specify a background image in our case we download this from unsplash and everything else you can leave it as is and then you could click on runtime run all and then it should work so now let's walk through this step by step so first let's run this cell and Define our image names then we want your L retrieve to download the background image from unsplash and we also need pillow image and matplotlib to load and plot the images then we apply your L retrieve this will download and also save this in a file so if we run this then it will plug the image here and if we refresh our folder then here we also have the background saved in this folder and then to develop our model we need opencv numpy torch and torch Vision they are all pre-installed in the collab so we don't have to worry about installing and then we Define a few helper functions first want to load the model and here we call torch upload and we can take this exact same line from the deep lip version 3 documentation that we can find at pytorch or slash Hub and then this model here so here we Define the vision and the version then here our model name and then we say pre-trained equals true and then here we find a few more information for example the mean and the standard deviation that was applied so we need to remember those values because we need to apply the same values in our pre-processing steps and then here one more thing I want to note is that we could also try different model structures like here we use the resnet 50 but we could also try a resonant 101 or the mobile net so yeah let's define our model and then we say model eval well because we don't need to train it we only need it for inference and then we return it and then we create another helper function to make a transparent foreground this gets the original image and then the mask that gets returned from this model so here we split this into all the different color channels blue green and red so note that opencv has it the other way around not red green and blue then we also add an alpha Channel and then we merge the alpha Channel back into this image so now we have the picture with an alpha Channel then we want to create a transparent background by first initializing it with zeros and with the same shape and then we set up the new mask with numpy stack and then we copy only the foreground color pixels from the original image where the mask is set and we can do this with numpy where so where the mask is set we use the picture pixels with the alpha and otherwise we use the zero background pixels and then we convert it to this type and return this and then we create the next helper function remove background so this will apply the whole pipeline this gets the model and the input file name and then here we first open the image with image from pillow then here we Define our pre-processing steps so first we convert the image to a tensor and then we normalize it with the same mean and standard deviation values that we found in the documentation here and then we want to apply the pre-processing steps and call unsqueeze to have this in the correct shape then we can also check if we have a GPU so inside this collapse you should you can move the model and the batch to the Cuda device and then we can call the model with the input batch and get the output and then we call Arc Max to have the

output predictions and then we want to create a binary so black and white mask of the profile foreground by saying the mask is the output prediction Spike dot cpu. numpi so we move this back to the CPU and then convert this to a numpy array again then we create our Background by initial analyzing this with zeros and then we say the mask is numpy where the mask is set we use 255 and otherwise we use the zero background pixels from here and then we apply our make transparent foreground helper function with the input image and this binary mask and this is our foreground and from this helper function we return both the foreground and The Mask so now let's run this cell and then here we set up everything so first we call load model and then we say remove background with the model and the input file so this should do all the processing now all right then this worked so let's plot our mask and let's see if this worked so yeah there we have it then let's also plot the foreground so yeah this works as well so the background pixels are now removed and now let's save this to a file so first we can say image from array and then the foreground and then we also if this is a JPEG file then we want to convert this to RGB if this is a for example a PNG file then we don't have to apply this and then we can say image save and now let's run this and now if we refresh the folder then now we also have the foreground saved here and now it's last step we want to merge this with a new background image so for this we Define a new helper function custom background this gets the background file and the foreground and this is currently a numpy array so we want to convert both to a pillow image for the array we say image from array and for the file we say image open and now we apply a crop box so here we Define the shape of the box and then we save background crop so now it essentially has the same shape as the foreground then we copy this and this will be our final image and now we put the foreground in the center of the background so we Define in our pastebox size like this and then we say final image paste so this is currently our background and into the background we paste the foreground into the box and the mask is also the final foreground and then we can return this so let's run the cell and then let's try it out so our final image is the custom background with the background file and this foreground array so let's run this and then let's also plot our image and yeah it worked so now as very last step we can also save it to our folder so let's refresh this and then we also have the final image here so yeah this worked as well and yeah that's it for this tutorial I hope you enjoyed this and if you did so then please leave us a like and consider subscribing to our Channel and then I hope to see you in the next video bye