Nails Semantic Segmentation iOS App Tutorial

Mobile developers are often asked to implement the latest available features for their platforms, and demand for ML models in production application has increased dramatically over the last couple of years.

The creation of production-ready Neural Networks requires a big dataset and lots of time, so our models in this course will have some reasonable limitations. But you will be able to train Semantic Segmentation Neural Network fast and understand critical concepts of how these models are trained and how they can be integrated into your apps.

Tutorial Overview

In this tutorial, we will look at the “Wanna Nails” case, and we will show you how to train a model that will detect nails in a couple of hours. “Wanna Nails” is an app that uses Object Segmentation to detect nails and try on different polish colors. You can download the app from here. It works like this:

We have trained a model using just 50 images. Here is an example of an app that you will get at the end of this tutorial:

Note: As we have mentioned above, our model has some limitations — it better detects the nails of a person whose nails we were using as a dataset. It will work worse in light conditions that weren’t presented in our dataset. Due to the lack of images in the dataset, it can sometimes make the wrong detections. But still, it’s a very good and fast way to make proof-of-concept and test app mechanics with Semantic Segmentation Neural Network that detects nails.

Dataset preparation

To create a production-ready dataset, we’ll need tens of thousands of different images in it. However, here we want to show you how you can train a Neural Network and integrate it with your app, and we hope that you’ll do it in a reasonable amount of time. Hopefully, a couple of hours. There are several big datasets on the web, such as COCO, PascalVOC, but they won’t give you any clues if you want to make something custom. So, our assumptions are the following — we need to prepare a dataset in a small amount of time, and we want it to give us consistent results from training.

Let’s imagine that our dataset needs some fuel. This fuel is images. To drive through the whole Route 66 successfully, we’ll need several tanks of gasoline. At this point, we don’t need this. Here we’ll try to make it from LA Downtown to Santa Monica with as little gallons as possible!

Here is a full dataset that we’ve used to train the nails segmentation model. It contains 50 images of nails. They are taken from different angles, from different backgrounds, with varying conditions of lighting. It allows the model to capture the same patterns in different situations.

We can apply the following rule of thumb for our datasets in this tutorial. Imagine that you have a baby, and this baby knows nothing about the environment, and you want to teach your baby to detect something. So, you want to provide her with enough examples to allow her to identify objects correctly. At this point, you can look at your dataset and think whether these examples would be enough for your baby to learn. Sounds a bit odd, but it helps :)

Images labelling using MakeML app

Let’s get back to work. So, you have captured 50 images of your hands, from different angles, from different backgrounds, maybe even in different lighting conditions. It takes me around 10 minutes to do it and to send it from my iPhone to my Mac. From this point, we need to markup them. To do this, we’ll use the MakeML app. You can download MakeML here, and learn more about it on their website.

Open MakeML app and create a new MakeML project.

Select Semantic Segmentation dataset type and Tensorflow training configuration, enter a project name and press the “Create” button.

Drag&Drop images that you’ve captured to Images Pane.

Create “nail” annotation title.

From this point, we’re ready to annotate our nails! You should annotate every nail as a separate object. To start annotation left-click on a border of a nail and circle it making clicks to create a shape. To end annotating a nail press “Return” button, and the app will connect the last and the first keypoints. To cancel the annotation, press the “Esc” button. To remove the last annotation point, press the “Delete” button.

It took me around an hour to annotate 50 images.

Semantic Segmentation model training

We also need to set training configuration parameters before training. For training nails model, I’ve used the following configuration:

Batch size — 8 Base learning rate — 0.002 Number of iterations — 2000

To start training at this point, all we need to do — is to press the “Run” button.

After you have pressed the “Run” button, the dataset uploads to MakeML servers and training of the model starts. Nevertheless, from now on, you need to wait until your model is ready, there is one parameter that you need to monitor. It’s Loss. You can see it in the output of your training.

Let me describe what this Loss means. Neural Networks usually trains at the following principle:

A model is making a prediction. It doesn’t know the correct answer. The correct answer is the annotations that we have made. A training algorithm compares the correct answer and the prediction of a model. In Semantic Segmentation, it looks for the overlapping of a predicted area and the exact area. Algorithm gets Loss from this overlapping. After this comparison of areas, the training algorithm tries to tune the model to minimize the error of the model.

So, Loss is an essential parameter that shows how big is the mistake of a model. It’s better when it’s less. You can ask, “What does it all mean? What should I look at?”. You need to compare the output of the training, and if Loss is decreasing — everything is fine. If it is going up, it means that something went wrong. It means that some of your model neurons have started to work wrong at some point.

We can recommend two actions to avoid this behavior. You can decrease the “base learning rate” (for example, from 0.007 to 0.004). You can reduce the “number of iterations” (for example, from 2000 to 1500), or you can make a combination of changes in these two training parameters.

TFLite model integration to iOS app

After the training of the model has finished, you can export it and receive the “.tflite” file ready for further use — press the “Export model” button.

The last step of our tutorial is the integration of the model into an iOS app. After exporting the model, you will receive a .zip file that contains a folder. We need a “result_model.tflite” file from this folder.

To make this tutorial concise, we have prepared a project and posted it on Github. All you need to do is to replace the “result_model.tflite” file with your model there.

To change colors that will be applied to your nails, you can change colors array in “ViewController.swift”.

Originally published at https://makeml.app.