Style Transfer on iOS

You can use the FritzVisionStyleModel to stylize images. Fritz provides a variety of options to configure predictions. These instructions will help you get Style Transfer running in your app in no time.

1. Build the FritzVisionStyleModel

To create the style model, you can either include the model in your bundle or download it over the air once the user installs your app.

Include painting styles in your app bundle

Fritz provides a pack of pre-trained style models, replicating styles of famous paintings.

Add the model to your Podfile

To use these style models, include the Fritz/VisionStyleModel/Paintings pod in your Podfile.

pod 'Fritz/VisionStyleModel/Paintings'

Make sure to install the recent addition.

pod install

Define FritzVisionStyleModel

Choose which FritzVisionStyleModel you want to use. You can choose a specific style or access a list of all models. There should only be one instance of the model that is shared across all predictions:

Using a specific style:

import Fritz
lazy var styleModel = FritzVisionStyleModel.starryNight

Using all styles:

import Fritz
let styleModels = FritzVisionStyleModel.allModels()

Using a specific style:

@import Fritz;
FritzVisionStyleModel *styleModel = [FritzVisionStyleModel starryNight];

Using all styles:

@import Fritz;
NSArray *styleModels = [FritzVisionStyleModel allModels];

Using custom style models

You can also create custom style models. See Create and manage custom style models for how to create your own style model.

Add FritzVision to your Podfile

To use these style models, include the Fritz/Vision pod in your Podfile.

pod 'Fritz/Vision'

Make sure to install the recent addition.

pod install

Define FritzVisionStyleModel

Assume you have a custom style model named CustomStyleModel:

import Fritz

let styleModel = FritzVisionStyleModel(model: CustomStyleModel())
@import Fritz;

FritzVisionStyleModel *styleModel = [[FritzVisionStyleModel alloc] initWithIdentifiedModel:[CustomStyleModel new]];

2. Create FritzVisionImage

FritzImage supports different image formats.

  • Using a CMSampleBuffer

    If you are using a CMSampleBuffer from the built-in camera, first create the FritzImage instance:

    let image = FritzVisionImage(buffer: sampleBuffer)
    
    FritzVisionImage *visionImage = [[FritzVisionImage alloc] initWithBuffer: sampleBuffer];
    // or
    FritzVisionImage *visionImage = [[FritzVisionImage alloc] initWithImage: uiImage];
    

    The image orientation data needs to be properly set for predictions to work. Use FritzImageMetadata to customize orientation for an image. By default, if you specify FritzVisionImageMetadata the orientation will be .right:

    image.metadata = FritzVisionImageMetadata()
    image.metadata?.orientation = .left
    
    // Add metdata
    visionImage.metadata = [FritzVisionImageMetadata new];
    visionImage.metadata.orientation = FritzImageOrientationLeft;
    

    Note

    Data passed in from the camera will generally need the orientation set. When using a CMSampleBuffer to create a FritzImage the orientation will change depending on which camera and device orientation you are using.

    When using the back camera in the portrait Device Orientation, the orientation should be .right (the default if you specify FritzVisionImageMetadata on the image). When using the front facing camera in portrait Device Orientation, the orientation should be .left.

    You can initialize the FritzImageOrientation with the AVCaptureConnection to infer orientation (if the Device Orientation is portrait):

    func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
        ...
        image.metadata = FritzVisionImageMetadata()
        image.metadata?.orientation = FritzImageOrientation(connection)
        ...
    }
    
  • Using a UIImage

    If you are using a UIImage, create the FritzVision instance:

    let image = FritzVisionImage(image: uiImage)
    

    The image orientation data needs to be properly set for predictions to work. Use FritzImageMetadata to customize orientation for an image:

    image.metadata = FritzVisionImageMetadata()
    image.metadata?.orientation = .right
    

    Note

    UIImage can have associated UIImageOrientation data (for example when capturing a photo from the camera). To make sure the model is correctly handling the orientation data, initialize the FritzImageOrientation with the image’s image orientation:

    image.metadata?.orientation = FritzImageOrientation(image.imageOrientation)
    

3. Run style transfer

Stylize images

styleModel.predict(image) { stylizedBuffer, error in
  guard error == nil, let stylizedBuffer = stylizedBuffer else { return }

  // Code to work with stylized image here.
  // If you're not sure how to use the output image, check out the public heartbeat
  // project.
}
[styleModel predict:image completion:^(CVPixelBufferRef result, NSError *error) {
  // Code to work with stylized image here.
  // If you're not sure how to use the output image, check out the public heartbeat
  // project.
}];

Configure style transfer model

Before running style transfer, you can configure the prediction with a FritzVisionStyleModelOptions object.

Settings
imageCropAndScaleOption

.scaleFit (default)

Crop and Scale option for how to resize and crop the image for the model

For example, you can build a model that will crop the image to the input size (640x480):

let options = FritzVisionStyleModelOptions()
options.imageCropAndScaleOption = .centerCrop

// Use the options in model prediction.
visionModel.predict(image, options: options) { result, error in

}
FritzVisionStyleModelOptions * options = [FritzVisionStyleModelOptions new];
options.imageCropAndScaleOption = FritzVisionCropAndScaleCenterCrop;

Create and manage custom style models

Train your own style model using the Fritz Style Training Template

If you want to train and use your own custom style transfer model, follow the open source training template on Google Colab.

For a full tutorial on training your custom model, take a look at this blog post on Heartbeat.

After you’ve finished training your model, follow these steps to add it to your iOS app.

1. Create a custom model in the webapp and add to your Xcode project.

For instructions on how to do this, see Core ML

2. Update your podfile:

pod 'Fritz/Vision'
pod install

3. Conform your model

For the following instructions, assume your model is named CustomStyleModel.mlmodel.

Create a new file called CustomStyleModel+Fritz.swift and conform your class like so:

import Fritz

extension CustomStyleModel: SwiftIdentifiedModel {

    static let modelIdentifier = "model-id-abcde"

    static let packagedModelVersion = 1
}

Create a new extension called CustomStyleModel+Fritz and conform your class like so:

// CustomStyleModel+Fritz.h
@import Fritz;

@interface CustomStyleModel(Fritz) <FritzObjcIdentifiedModel>
@end
// CustomStyleModel+Fritz.m
#import "CustomStyleModel+Fritz.h"

@implementation CustomStyleModel(Fritz)
+ (NSString * _Nonnull)fritzModelIdentifier {
    return @"model-id-abcde";
}

+ (NSInteger)fritzPackagedModelVersion {
    return 1;
}
@end

Fetch custom style models by tag

If you create your own styles using the Fritz Style Training Template, you can Add Tags to those models and dynamically load them in the SDK.

When loading style models by tag, you do not need to include them in your bundle:

import Fritz

FritzVisionStyleModel.fetchStyleModelsForTags(tags: ["my-custom-style-models", "impressionists"]) { models, error in
    guard let styleModels = models, error == nil else {
        return
    }
    // Instantiated models live here.
}

Note

If you have many Style Models for a given tag query, loading all models into memory could cause memory pressure. In this case, you can load the models on demand.

1. Get FritzManagedModels from ModelTagManager

For detailed instructions on how to get the list of FritzManagedModel, refer to Querying Tags in the iOS SDK.

2. Instantiate FritzStyleModel from FritzManagedModel

You should get a FritzManagedModel list containing all models matching your tag query. From those models, you can instantiate a single instance of the style model as you need.

import Fritz

var styleModel: FritzVisionStyleModel?

// Get managed models for tags that have been previously loaded..
let tagManager = ModelTagManager(["style-transfer", "custom-models"])
let managedModels = tagManager.getManagedModelsforTags()

for model in managedModels {
    managedModel.fetchModel { mlmodel, error in
        guard let fritzMLModel = mlmodel, error == nil else {
            return
        }
        try {
            styleModel = try FritzVisionStyleModel(model: fritzMLModel)
        } catch {
            // Errors may occur if the style model does not contain the appropriate inputs and outputs.
        }
}

In this example, the single Style Model is replaced by each model as it iterates through the for loop, but you could have a user choose from a list of options returned from the tag query.

Access model metadata

Metadata is easily accessible on the FritzStyleModel object. To add metadata to models, read Adding Metadata to Models:

import Fritz

// Assume you have an instantiated Style Transfer model
let styleModel = FritzVisionStyleModel(model: ...)
let metadata: [String:String]? = styleModel.metadata

Additional style model resources