iOS Image Segmentation

Note

If you haven’t set up the SDK yet, make sure to go through those directions first. You’ll need to add the Core library to the app before using the specific feature API or custom model. Follow iOS setup or Android setup directions.

You can use a Fritz Image Segmentation Model to partition an image into multiple segments that recognize everyday objects. These instructions will help you get Image Segmentation running in your app in no time.

1. Build the FritzVisionPeopleSegmentationModel

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

Include pre-trained image segmentation models in your app bundle

Add the model to your Podfile

Include the segmentation model in your Podfile. You can choose between three different models:

We offer models at two different resolutions to segment people, 384x384 and 768x768. The smaller model will run faster while the larger model will produce cleaner masks but is slower for real time cases.

People 384x384

pod 'Fritz/VisionSegmentationModel/People'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

People 768x768

pod 'Fritz/VisionSegmentationModel/PeopleMedium'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

Segment pixels that contain hair.

pod 'Fritz/VisionSegmentationModel/Hair'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

A two-class model that identifies pixels that contain either people or pets. Produces masks that are 768x768.

pod 'Fritz/VisionSegmentationModel/PeopleAndPetMedium'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

Segment pixels that contain pets. Model output size is 224x224.

pod 'Fritz/VisionSegmentationModel/Pet'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

Segment pixels that contain the sky. Model output size is 224x224.

pod 'Fritz/VisionSegmentationModel/Sky'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

To idenfity the following object in your living room (with the colors that represent each segment in the final result):

  • Chair (Sandy Brown)
  • Wall (White)
  • Coffee Table (Brown)
  • Ceiling (Light Gray)
  • Floor (Dark Gray)
  • Bed (Light Blue)
  • Lamp (Yellow)
  • Sofa (Red)
  • Window (Cyan)
  • Pillow (Beige)
pod 'Fritz/VisionSegmentationModel/LivingRoom'

Make sure to install the added pod:

# Update the podspec repo with the latest version of Fritz
pod repo update

pod install

To idenfity the following outdoor objects:

  • Building (Green)
  • Sky (Yellow)
  • Tree (Blue)
  • Sidewalk (Orange)
  • Ground (Purple)
  • Car (Teal)
  • Water (Purple)
  • House (Sandy Yellow)
  • Fence (Light Pink)
  • Sign (Dark Blue-Green)
  • Skyscraper (Purple-Pink)
  • Bridge (Chocolate)
  • River (Light Yellow)
  • Bus (Dark Red)
  • Truck (Light Green)
  • Van (Moss Green)
  • Motorbike (Beige)
  • Bicycle (Dark Blue)
  • Traffic Light (Grey)
  • Person (White)
pod 'Fritz/VisionSegmentationModel/Outdoor'

Make sure to install the added pod:

pod install

Define FritzVisionPeopleSegmentationModel

Choose which segmentation model you want to use. There should only be one instance of the model that is shared across all predictions. Here is an example using the FritzVisionPeopleSegmentationModel:

import Fritz

let peopleModel = FritzVisionPeopleSegmentationModel()
@import Fritz;

FritzVisionPeopleSegmentationModel *model = [[FritzVisionPeopleSegmentationModel alloc] init];

Note

Model initialization

It’s important to intialize one instance of the model so you are not loading the entire model into memory on each model execution. Usually this is a property on a ViewController. When loading the model in a ViewController, the following ways are recommended:

Lazy-load the model

By lazy-loading model, you won’t load the model until the first prediction. This has the benefit of not prematurely loading the model, but it may make the first prediction take slghtly longer.

class MyViewController: UIViewController {
  lazy var model = FritzVisionPoseModel()
}

Load model in viewDidLoad

By loading the model in viewDidLoad, you’ll ensure that you’re not loading the model before the view controller is loaded. The model will be ready to go for the first prediction.

class MyViewController: UIViewController {
  let model: FritzVisionPoseModel!

  override func viewDidAppear(_ animated: Bool) {
    super.viewDidAppear(animated)
    model = FritzVisionPoseModel()
  }
}

Alternatively, you can initialize the model property directly. However, if the ViewController is instantiated by a Storyboard and is the Initial View Controller, the properties will be initialized before the appDelegate function is called. This can cause the app to crash if the model is loaded before FritzCore.configure() is called.

Download the model over the air

Add FritzVision to your Podfile

Include Fritz/Vision in your Podfile.

pod 'Fritz/Vision'

Make sure to run a pod install with the latest changes.

pod install

Download Model

import Fritz

var peopleModel: FritzVisionPeopleSegmentationModel?

FritzVisionPeopleSegmentationModel.fetchModel { model, error in
   guard let downloadedModel = model, error == nil else { return }

   peopleModel = downloadedModel
}
@import Fritz;

FritzVisionPeopleSegmentationModel* model;

[FritzVisionPeopleSegmentationModel fetchModelWithCompletionHandler:^(FritzVisionPeopleSegmentationModel * _Nullable result, NSError * _Nullable error) {
    model = result;
}];

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 image segmentation model

Run image segmentation

Run image segmentation on input image:

guard let result = try? peopleModel.predict(image) else { return }
[model predict:image options:options completion:^(FritzVisionSegmentationResult * _Nullable result, NSError * _Nullable error) {
    // Use segmentation result.
}];

Configure Image Segmentation Model

Before running image segmentation, you can configure the prediction with a FritzVisionSegmentationModelOptions 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 (384x384):

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

guard let result = try? peopleModel.predict(image, options: options) else { return }
FritzVisionSegmentationModelOptions * options = [FritzVisionSegmentationModelOptions new];
options.imageCropAndScaleOption = FritzVisionCropAndScaleCenterCrop;

[peopleModel predict:image options:options completion:^(FritzVisionSegmentationResult *result, NSError *error) {
    // ...
}];

4. Create masks from image segmentation result

Create image with all identified classes

Create an image with each class as a different color.

let result: FritzVisionSegmentationResult = // ... Result from prediction above

let peopleMask = result.buildMultiClassMask()

Pass in a minimum threshold to only display pixels with a confidence score at least as high as the withMinimumAcceptedScore

let result: FritzVisionSegmentationResult = // ... Result from prediction above

let peopleMask = result.buildMultiClassMask(withMinimumAcceptedScore: 0.7)

Mask a specific class

You can create a mask for a specific ModelSegmentationClass:

let result: FritzVisionSegmentationResult = // ... Result from prediction above

let peopleMask = result.buildSingleClassMask(forClass: FritzVisionPeopleClass.person)

Additionally, you can call buildSingleClassMask with clippingScoresAbove and zeroingScoresBelow arguments as well. This is helpful for dealing with the uncertainty of the model.

  • When clippingScoresAbove is set, any pixels with a confidence score above that threshold will have an alpha value of 255 (completely opaque).
  • When zeroingScoresBelow is set, any confidence scores below will not appear in the mask.
  • Any scores between clippingScoresAbove and zeroingScoresBelow will have a value of classProbability * 255. It’s useful to create a blur around predictions that may still contain the desired class.
let result: FritzVisionSegmentationResult = // ... Result from prediction above

let peopleMask = result.buildSingleClassMask(
    forClass: FritzVisionPeopleClass.person,
    clippingScoresAbove: 0.7,
    zeroingScoresBelow: 0.3)

Change Hair Color with the Hair Segmentation Model

The hair segmentation model can be used to blend the color of the mask with the original image.

1. Run prediction on FritzVisionHairSegmentationModel

let model = FritzVisionHairSegmentationModel()

let image = FritzVisionImage(...) // ... Image from above

guard let result = try? model.predict(image) else { return }

2. Create mask from prediction result

Here you can change the senstivity of the hair mask. In this example, any pixels with a confidence score above 0.6 will be completely opaque (alpha of 255) in the resulting mask. Values below 0.2 will completely transparent.

guard let mask = result.buildSingleClassMask(
  forClass: FritzVisionHairSegmentationClass.hair,
  clippingScoresAbove: 0.6,
  zeroingScoresBelow: 0.2,
  resize: false,
  color: .blue) else { return }

3. Blend mask with original image

Here you can blend the mask with the input image. It will scale the mask to cover the original input image.

let blended = image.blend(
  withMask: mask,
  blendMode: .softLight,
  interpolationQuality: .medium,
  opacity: 0.7
)

Create and use a custom segmentation model

Train a custom segmentation model

If you want to train and use your own custom image segmentation model, follow the direction in the Fritz Image Segmentation Repo.

Use a custom segmentation model

In this tutorial, imagine you are training an animal segmentation model on 4 different classes: None (matching no ani mals), Bird, Dog, and Horse. In your trained model, they correspond indices to 0, 1, 2, 3 respectively.

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

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

For instructions on how to do this, see Integrating a Custom Core ML Model.

2. Conform your model

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

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

import Fritz

extension CustomAnimalSegmentationModel: SwiftIdentifiedModel {

    static let modelIdentifier = "model-id-abcde"

    static let packagedModelVersion = 1
}

3. Define the AnimalSegmentation Model

First define the classes used in your model. We recommend you create a class to hold all individual classes. Each ModelSegmentationClass is created with a name and a color used for the mask.

import Fritz

public class AnimalClass: NSObject {
    public static let none = ModelSegmentationClass(label: "None", index: 0, color: (0, 0, 0, 0))
    public static let bird = ModelSegmentationClass(label: "Bird", index: 1, color: (0, 0, 0, 255))
    public static let dog = ModelSegmentationClass(label: "Dog", index: 2, color: (0, 0, 128, 255))
    public static let horse = ModelSegmentationClass(label: "Horse", index: 3, color: (230, 25, 75, 255))

    public static let allClasses: [ModelSegmentationClass] = [.none, .bird, .dog, .horse]

}

Next, create the image segmentation model:

let animalSegmentationModel = FritzVisionSegmentationModel(model: CustomAnimalSegmentationModel(),
                                                           name: "AnimalSegmentationModel",
                                                           classes: AnimalClass.allClasses)

4. Create a FritzImage and segment your image

Follow steps 3. Run image segmentation model and 4. Create masks from image segmentation result to use your custom segmentation model.

Additional Resources

Want more resources to get started? Check out the following: