Custom Face Mask Segmentation for Lens Studio with Fritz AI (Video + Transcript)

Building and implementing custom machine learning model with SnapML in Lens Studio—no code required

Follow along as we build a custom, ML-powered face mask color changing Snapchat Lens.

You’ll learn how to use Fritz AI Studio to:

🧑‍💻 Generate a new image dataset

😷 Train a Face Mask Segmentation model without code

✨ Implement in Lens Studio via a pre-configured project template.

For more on working with Fritz AI & SnapML in Lens Studio, check out the following resource:

👉 SnapML Collected Resources: Our team’s collection of explainers, tutorials, spotlights, and more, covering all things SnapML and Lens Studio

We’ve included the transcript for this video below, as well.

Hi there — My name’s Austin, and I’m the Head of Content and Community at Fritz AI. In this tutorial, I’m going to walk you through the steps of training a custom machine learning model using Fritz AI, and then implementing that trained model inside Lens Studio.

Using custom ML inside Lens Studio — with the help of SnapML — allows Lens Creators to extend the capabilities of Lens Studio, applying unique artistic effects, recognizing and tracking custom objects (like brand logos or product packaging), and create unique computer vision triggers for AR effects.

Intro Project

An essential step of any ML project, before you start building, is defining what you want your model — and thus your Lens — to do in the real world.

For this tutorial, we want to create a Lens that allows us to change the color of a face mask/covering, but leave the rest of the scene unchanged. Like this (show final result).

This means we want to segment out face masks present in a given scene, so that we can apply a color changing effect to them while leaving the rest of the scene unchanged. Additionally, we want our ML model to perform this segmentation when masks are being worn by users — not when the masks are sitting on a table, hanging from an ear, etc. This context will help when we create our dataset. More on that in a moment.

Create Fritz AI Project

To get started, the first step is to create a new project in your Fritz AI Studio account. We can create a new project using the left dropdown menu on the Fritz AI homepage. I’ll enter a simple name for our project, and for our project type, we’ll choose Custom Trained.

Fritz AI Studio currently supports three types of custom models for SnapML/Lens Studio: Image Labeling, Object Detection, and Image Segmentation. For this tutorial, we’re going to work with Image Segmentation

When you create a project, you can also choose to upload some dataset images to get started quickly. For this tutorial, I’m going to skip this option so I can show you the full process of creating a dataset.

Fritz AI Studio Overview

Before we jump in, a short overview of Fritz AI Studio. The Studio is our no-code machine learning platform that empowers you and your team to build completely custom ML models that are ready-to-use inside Lens Studio. You can do things like track custom objects, isolate and alter scene elements, provide unique interactive triggers for AR effects, and much more.

Almost everything you need to build your models is available inside Fritz AI Studio: with the platform you can generate and collect labeled datasets, train Lens Studio-optimized models without code, and then deploy them directly to Lens Studio, either as model files or Lens Studio template projects.

We’re going to work through this entire process, end-to end. Specifically we’ll:

  • Build an initial dataset
  • Train our first model
  • Implement and test that model inside Lens Studio via a project template

Adding an Image Collection

Our first step, now that we’ve defined the problem and created a new project, is to collect and label some data — the life blood of machine learning. When starting a project from scratch, building an initial dataset can be a daunting task. But fear not. Rather than collecting and labeling thousands of images of face masks by hand, we’ll use our Dataset Generator to programmatically do some of this work for us.

All we’ll need to do is create and manually label a handful of what we call “seed” images. We’ll learn more about these in a moment.

Creating a new project automatically creates what’s called a new seed image collection for us. Image collections contain both images and their corresponding annotations, which will eventually be used to train our model.

About Seed Images

Now it’s time to upload some seed images. First, some important info about seed images.

Seed images are transparent PNGs — they’ll be used like stickers that will be programmatically pasted onto random backgrounds. This will allow us to generate a large, diverse set of images for model training.

Though project requirements can vary, we recommend you start with 20–50 seed images for each type and class of object you want to predict. So for this face mask segmentation example, that would mean uploading 20–50 images of face masks, all with their backgrounds removed.

As another example, if you wanted to train an image labeling model that distinguishes between three different types of food items, you’d want to upload and label 60–150 seed images — 20–50 for each different food item (each class).

We have to make a couple decisions here. For our model, we could make seed images of just the face mask itself, but including the subject’s head and shoulders in some of the seed images might result in better model performance — especially since our intention is to segment masks while people are wearing them.

As a consequence, the model will learn what a face mask looks like in the context of a human face/head — so a mask sitting on a table, for example, would be less likely to be detected.

Additionally, the greater the variety in your seed images — angle, rotation, images with and without facial hair, etc. — the more accurate your model will be in the real world.

There are quite a few free and easy-to-use tools out there for creating these transparent PNGs. Tools like remove.bg, Photoshop, and even Preview on MacOS can do this.

I prefer remove.bg, but feel free to use the software that feels most natural to you. Image databases like Google Images, Unsplash, Pixabay, and others are good places to start when looking for seed images, but be sure any photos have appropriate licensing.

Uploading Seed Images

Remember, we want between 20–50 labeled seed images for our first dataset. For demonstration purposes, I’m only going to upload and label 3 right now, but the full dataset for this project would need to include more seed images. We can easily bulk upload those — I prefer to keep all my seed images in a separate folder, just so they’re easier to find.

Now that we’ve uploaded our seed images, we need to label them. We can do this right inside Fritz AI, so no need to use an external labeling tool. Inside the image collection, we can click the ANNOTATE IMAGES button to open the annotation tool.

Annotating Seed Images

The first step is to create an object or set of objects that can be annotated in each image. When we we start annotating images in a new collection, you can see that we’re asked to define these objects upon launching the annotation tool.

This set of objects defines what your model will be able to predict. Here, we only have one object: Face Mask. But you can easily create more objects to be annotated if needed.

Once this object is defined and confirmed, we can start annotating these seed images. Remember, I’m only going to work through 3 of these, just as an example, but you’ll want your datasets to include quite a few more seed images.

If you’re new to the labeler, it takes a just a little getting-used-to — I’d personally recommend familiarizing yourself with and using the hotkeys or keyboard shortcuts. So to label our first image, we select the face mask object, and then use the drawing tool to draw a segmentation mask around the…well, face mask.

We’ll do this with each image, being sure to save each annotation once we’re happy with it before moving on to the next. You can also always delete annotations you’re unhappy with and start over. I’ll speed up the process here, but take care to ensure your seed image labels are as accurate as possible.

Generating a Snapshot

After all seed images have been annotated, you can use the image collection to generate what’s called a dataset snapshot. Snapshots are sets of images and their annotations frozen in time and are used to train models.

From the Datasets tab in the left hand nav, click the CREATE A NEW SNAPSHOT button. Provide a name, ensure you have the correct seed image collection selected, and choose the number of images you want to generate. We recommend 50–100 snapshot images per seed image (i.e. for 50 seed images, a snapshot of 2500–5000 images).

As a general rule, more good data is better. Again, for demonstration purposes, I’m just going to generate a snapshot of 150 images.

As you can see on the snapshot creation form, you can also use a couple advanced configuration options, one of which is adjusting the foreground scale range, which will control the relative size of the seed image stickers on each snapshot image. each image.

The goal is to generate images that contain objects that most closely approximate what the model might see in real world use.

One important thing to keep in mind…While image collections can be changed over time (e.g. adding new images or modifying annotations), snapshots are fixed. Once you’ve created a snapshot, it will never change. If you make a change to an image collection, you will need to create a new snapshot in order to reflect those changes. This way, snapshots provide an audit trail, making sure the exact data used to train a model is always available.

Previewing Snapshot Images

Generating your dataset should only take a few minutes, but while you wait, you can check out this dataset preview to get a better sense of how the Dataset Generator actually works, and to see if you’ve picked an appropriate foreground scale range. See how the seed images we’ve labeled are pasted on a whole bunch of random backgrounds? That’s the result of what we call synthetic data generation.

Training a Model

Now we’ve successfully generated a snapshot, we’re ready to train a first version of our model! Model training is really simple in Fritz AI — no code or Python notebooks or any other model-building overhead. To do this, click on the Training tab on the left hand menu, and then “TRAIN A MODEL” over on the right

All we need to do is provide a name, optional description, choose the snapshot (or snapshots) we want to train the model on, and then set a training budget.

We recommend 5 training hours for custom models, but do note that if your model stops improving during training, Fritz AI will automatically stop the training process, and you won’t be charged for hours that are leftover.

So here I’ll keep the suggested number of hours, makes sure that I’ve chosen the correct snapshot, and then hit train. That’s it! Our model is now training — similar to the dataset snapshot we generated, we’ll get an email confirming when it’s trained

I’m going to grab a bit of food while we wait — See you in just a bit.

MODEL TRAINING INTERLUDE — ELEVATOR MUSIC

Download LS Project

Hey again!

So I just got an email letting me know my model has finished training successfully — to show you how this works in Lens Studio, I actually ended up training a different version of the same model with more seed images and a bigger dataset snapshot. But it’ll still do the same intended task, just hopefully with higher accuracy.

If we jump back into our Fritz AI account, we can find the trained model under the “Models” tab in the lefthand menu — if you want to see more details about your model training job, visit the “Training” tab.

Fritz AI gives us two options for accessing our trained model. We can download the model file directly and add it to Lens Studio manually, or we can download a zip file that includes a ready-made Lens Studio project file and a couple other resources.

For Creators just getting started with ML, we recommend the project template, which is what we’ll work with here.

All right so I’ve downloaded and unzipped the file. A couple of things to pay attention to. First, the .lsproj file right here. This will launch Lens Studio and initialize a template project. Second, this readme file inside the “Public” folder — we’re going to work through these steps in a moment, but this readme also provides these step-by-step instructions. And third, we have a “Previews” folder, which includes some image and video assets for testing the Lens in Lens Studio.

Keep in mind that the following steps are for image segmentation models — other model types require slightly different wiring inside Lens Studio, so be sure to reference the readme in your project file. Additionally, Fritz AI segmentation project templates come pre-loaded with a color changing UI, as a way to show how models function inside Lens Studio.

Import your model

Upon opening the Lens Studio project file, you will be shown a “Model Import” dialog box. You do not need to change any of these settings. Just click “Import”. You might also see a dialogue box pop up asking if you want to compress assets. Click “Yes”.

Add your model to the ML Component

Fritz AI projects come with an ML Component (or ML Controller, in some cases, where more scripting is needed) already setup. We just need to add our model to the Component, and do a bit of input and output wiring

If we click on the ML Component in the Objects panel, its configuration will open up in the inspector. In the Component’s “Model” field, we first add our model, which should be in the resources panel.

Next, we need to set the model’s input texture to “Device Camera Texture”. This will feed what the camera sees to the segmentation model.

We also need to create an Output Texture, which will allow us to use the predictions that the model outputs (i.e. predicting the pixels that belong to “Face Masks”). To do this, all we need to do is click “Create Output Texture

in the model output options. This texture will be used by other Lens Studio elements to access the output of your model.

Connect the model output texture to the Fritz Segmentation Texture material

We’re getting really close! Now we need to connect the predictions that the model is outputting to a segmentation material (again, pre-configured in the project template).

From the Resources list in the bottom left corner of Lens Studio, select the “Fritz Segmentation Texture” material. In the Inspector on the right, click the Model Output option under Graph Parameters and set the value to the output texture you just created.

Preview & Test

And that’s it! With that, we’ve successfully trained and implemented a face mask segmentation model in Lens Studio, resulting in a face mask color changer. Feel free to preview and test on the images provided in the project file, from your webcam, or push to Snapchat and test there!

Next Steps/Outro

We’re excited about ML inside Lens Studio for a wide variety of reasons, but one of the main reasons is because it actually extends what’s possible in Lens Studio, beyond functionality included in other template projects or tools.

Taking this custom model, we could experiment with different AR effects beyond color, use model predictions as a trigger for other unique effects, and much more.

If you’d like to experiment with a face mask segmentation model, we have one already created and available in our pre-trained model zoo. We’d love to see what you come up with.

Be sure to subscribe for more tutorials and content from our team, and check out some helpful link in the description below. Until next time!

Avatar photo

Fritz

Our team has been at the forefront of Artificial Intelligence and Machine Learning research for more than 15 years and we're using our collective intelligence to help others learn, understand and grow using these new technologies in ethical and sustainable ways.

Comments 0 Responses

Leave a Reply

Your email address will not be published. Required fields are marked *