CreateML Tutorial for iOS: Creating a Customized Image Filter using Style Transfer

Apple’s machine learning frameworks have become more and more powerful over the years. Style Transfer, or Neural Style Transfer, one technique which the frameworks provide. Through style transfer, you turn one image into a new image, having it stylized as a third. In this tutorial, you’ll learn how to leverage Apple’s ML Frameworks to achieve this directly on your iOS device!

Machine learning (ML) in iOS and macOS relies on two core Apple frameworks — CreateML and CoreML. While CreateML allows you to build & train an ML model, CoreML allows you to run these models.

CreateML initially was only available for macOS for generating models. You can then bundle these models for use in macOS, iOS or iPadOS applications.

In 2021, Apple made this framework available for iOS 15 and iPadOS 15. Now, CreateML framework allows on-device model training and creation.

In this tutorial, you’ll learn:

• How to create an ML model on-device.
• How to use the model to generate a stylized image.
• Core concepts in ML training and prediction.

Getting Started

Use the Download Materials button at the top or bottom of this tutorial to download the starter project.

The starter project is a simple app called Petra.

Open the starter project in Xcode, then build and run it on your device. Make yourself familiar with the UI of the app and the codebase.

This app has a single screen that allows you to choose two images. The first image is a style image which indicates what style is desired. The second image, the user-selected image, will have the desired style applied to it.

The button on the bottom Apply Style! should transform the user image into the selected style. This button will work after you put in the missing pieces by the end of this tutorial.

For this app, you’ll create an MLModel that you’ll use to stylize photos of your pet.

In this tutorial, you’ll train the model to apply the style to a pet such as a cat or a dog — or an elephant or a chicken. You can apply the same principles to create your model trained just for your pet.

What is Style Transfer?

Style transfer is a technique that blends the content of one image with the style of another. An ML model learns the style, including pattern, color or texture, from a style input. It then uses it to recompose the content input in that style.

The model creates building blocks from the style content. It then recreates your content with these blocks, just as you’d do in a game of Minecraft.

In the example above, a tiled/pattern image A acts as a source for extracting the style (e.g., color, pattern, texture). The style is applied to the image B to generate a stylized output image C.

Sounds like magic, doesn’t it?

This technique has varied applications ranging from digital art to virtual reality. It has also found its place in gaming as demoed in the Google Stadia platform. Apple’s ARKit also supports style transfer to stylize a specific scene.

Preparing the Dataset

To create a model, you must work with a training set using a machine learning algorithm. The CreateML framework in iOS is no different and provides all the functionality you need to train a model.

The training data should be as close as possible to the expected user content.
For instance, for an app that stylizes a human face, you’ll need portraits of people to train the model.

For this tutorial, you’ll need images of pets for training the model. You’ll use a tiny subset of this data from Kaggle containing cat and dog pictures. Don’t worry about downloading the entire dataset. The sample project already has our small dataset, which you’ll use for the tutorial.

Open the project directory in Finder and find the TrainingData folder.

For demo purposes, the project contains a sample stylized image. You can see the PresetStyle_1 image under Assets.xcassets in the project.

To recap, this is what you have in the starter project:

1. Training Data: A directory containing pet images as our training set.
2. Sample Style Image: A sample artistic or tiled image with a distinct texture, color or pattern.

Next, you’ll create and train a model using the data.

Training the Model

As with any machine learning, model training is the primary step in style transfer. Let the training begin!

Open MLModelTrainer.swift in the MLModelHandling group. Replace the contents of
trainModel(using:validationImage:sessionDir:onCompletion:) method with the code below.

// 1
let dataSource = MLStyleTransfer.DataSource.images(
    styleImage: styleImage,
    contentDirectory: Constants.Path.trainingImagesDir ?? Bundle.main.bundleURL,
    processingOption: nil)
// 2
let sessionParams = MLTrainingSessionParameters(
    sessionDirectory: sessionDir,
    reportInterval: Constants.MLSession.reportInterval,
    checkpointInterval: Constants.MLSession.checkpointInterval,
    iterations: Constants.MLSession.iterations)
// 3
let modelParams = MLStyleTransfer.ModelParameters(
    algorithm: .cnn,
    validation: .content(validationImage),
    maxIterations: Constants.MLModelParam.maxIterations,
    textelDensity: Constants.MLModelParam.styleDensity,
    styleStrength: Constants.MLModelParam.styleStrength)

Here’s what’s going on with the code:

1. You create dataSource which is a MLStyleTransfer.DataSource. This tells the model where the data is to train from.
2. You specify parameters to the training session through MLTrainingSessionParameters:
   a. sessionDirectory: To hold the data generated during the session.
   b. reportInterval: Number of iterations after which the session reports the progress.
   c. checkpointInterval: Number of iterations after which the session saves a checkpoint.
   d. iterations: Total number of iterations to run for the session.
3. You also set ModelParameters:
   a. algorithm: The style transfer task’s training algorithm that prioritizes either speed (.cnnLite) or quality (.cnn).
   b. validation: Image used for validation during the training — one from the training set.
   c. maxIterations: Maximum number of training iterations the style transfer model can use. This is ignored when you set the iterations in sessionParams.
   d. textelDensity: Defines how coarse or fine the style details are.
   e. styleStrength: Specifies how heavy the influence of style will be on the content image.

You’ll review the effects of some of these parameters in the following sections.

Finally, add the following code after modelParams in the same method:

// 4
guard let job = try? MLStyleTransfer.train(
    trainingData: dataSource,
    parameters: modelParams,
    sessionParameters: sessionParams) else {
    onCompletion(nil)
    return
}

The MLStyleTransfer.train(trainingData:parameters:sessionParameters:) is one method that does it all! You pass the above created parameters and the data source to it. It returns an MLJob object for handling the training result.

Calling this method starts an asynchronous style transfer model-training session.

Next, you’ll learn how to use the output of the training.