[Guide] Style Transfer using SnapML

Hello, this to my tutorial about working with SnapML in order to create your own Style Transfer lens within Lens Studio.

Initially before starting to make Style Transfer lenses, I found some scattered tutorials on YouTube and the Lens Studio forum, however I’d like to consolidate it all into one post of what I’ve learned. Sorry if I’m repeating what others have said before, but if this helps anyone out I’d be happy.

First, in order to use the Style Transfer template, we need a model. In order to obtain a model with our desired image output, we need to train our machine learning model using Google Colab.

If you don’t know what Colab is, it is a (mostly- there’s a paid pro version if desired) free service that allows you to run programs within a collaborative setting; comparable to a Google Docs, but with Python instead of word documents. The most important part here is to download the SnapML Style Transfer notebook from the Lens Studio support page. This notebook contains modules of code, in order, that will allow us to specify our style image and our desired number of training steps to get our output image to come out as expected.

Before we can use the notebook, however, we need a style image. This image is inserted into our notebook and used as the desired outcome image, meaning the ML model will use Colab’s database of pictures to train the machine appropriately. Once you come up with an idea for your lens, you can look for images on Google Images or other image hosting sites that can be used for your model. For example, if you were trying to make the world appear as Skittles, you might use a picture of skittles arranged as, say, a person. This way, the ML model sees what it needs to train towards.

Once you find your style image, you can begin setting up for the training process. Download the example notebook Snapchat provides here (direct download) and upload it to Colab (it has a .ipynb extension).

There are a few important things you want to note when using the service. First, looking at the left-hand side under the CO logo, you’ll find a folder icon. If it’s not already opened yet, click on it. This is your File Explorer, and here you are able to upload any style or test images you need. To upload them, simply drag and drop them into that area. A test image is provided in the .zip file you downloaded earlier.

To use your images in the actual program, be sure to navigate to the section of code titled “Global variables for training”. Here, you should see a variable called STYLE_IMAGE_PATH. Within the single quotations on the other side of the equal sign, enter the name of your image along with the extension of the file. Make sure your extension matches!

The other variable you need to pay attention to is NUM_TRAINING_STEPS. This is the number of steps that the model will train for. The higher the number, the better your model will come out and the better your output image will be, however it will take longer during both the setup and the training process.

Finally, if you’d like the file to automatically download once the ONNX file is created, scroll all the way down to the end of the notebook and uncomment the function download_onnx_from_colab() by deleting the hashtags.

Now, navigate to the top bar and click on Runtime, then Run All. This will run each module of the notebook in order and create the ONNX file we will use in Lens Studio. It is important that you keep your computer running and connected to the internet until the file is downloaded because your files are deleted from your runtime after disconnecting or being idle for a period of time. If you scroll down to the Training loop module, you’ll be able to see the current output of the model every 250 training steps.

Once the process is complete, you will automatically download a style_new.onnx file. Open Lens Studio and create a new project using the Style Transfer template. Drag your ONNX file into the Resources panel in the bottom left portion of the screen. Then, select “ML Component” in the Objects panel in the top left portion of the screen, and drag your style_new resource into the “Model” section of the ML Component panel in the Inspector. You should now see your model reflected in the preview portion of the screen!

If you don’t like the current look of your model, you can change some settings by playing with the output image under Orthographic Camera > Style Transfer Output. The two big things you can play with are the Alpha, or transparency, of the model, and the blend mode of the model. This will help you create your desired result.

That’s basically all there is to know about the Style Transfer template using Lens Studio and Snap’s powerful SnapML. If you have any questions feel free to reach out!