Remove the image background in Swift using Core ML
For my personal photo editing indie App Cre8, I wanted to implement the functionality to make Avatars for social networks. For that user would have to add an image and erase all the background. I tried to do that myself a couple of times and found that process very long and annoying, so I had to think about something that will do that automagically. So, let’s try to use the help from Machine Learning to achieve expected results.
What is CoreML, and how to start using it?
Core ML is Apple’s framework to help integrate trained models into the app.
A model is the result of applying a machine learning algorithm to a set of training data. You use a model to make predictions based on new input data. Models can accomplish a wide variety of tasks that would be difficult or impractical to write in code. For example, you can train a model to categorize photos, or detect specific objects within a photo directly from its pixels.
We can train the model on a large dataset. If you have no great amount of data and any willingness to train the models, you can use existing Core ML Models provided by Apple.
For our purposes, we are going to use the DeeplabV3 model.
After you have downloaded it, drag a DeepLabV3.mlmodel file to your project in the Xcode.
Ok, that’s probably enough for the theory, let dive into the code.
All the work will be done on the UIImage extension, so be prepared for that:
Firstly, I would like to say thank you to Matthijs Hollemans for making CoreMLHelpers open source and available for everyone. All the image transformations will be done using those helpers.
- We have to initialize our model, and for that, we are going to use the generated Swift model class DeepLabV3.
And for that we can make a helper method to get the model:
2. Resize our image to match the model’s input image parameters. You can see that in Predictions input and output. If you have missed that, scroll up a bit and check the screenshots or explore DeepLabV3.mlmodel file.
3.1. Return the image without the background.
3.2. Return the background for drawing on CGContext.
CIImage extension for removing white pixels and composing images:
In conclusion, I would like to say that results could be very surprising in both ways. Many times it almost ideally clears the background, but for some images, where the foreground and background colors are pretty similar you might want to give the user additional tools like the eraser to achieve a better final result.
The final code is available here: Gist
The result could be seen in the Cre8 photo editor app.
Special thanks to my wife (who just started learning iOS development) for the support and permission to use her photos. 🤗 ❤️
