iOS Image Segmentation

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:

To identify people (segments that represent people are marked in black)

pod 'Fritz/VisionSegmentationModel/People'

Make sure to install the added pod:

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:

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];

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:

peopleModel.predict(image) { result, error in
    guard error == nil, let segmentationResult = result else { return }

    // Use segmentation result.
}
[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

peopleModel.predict(image, options: options) { results, error in
    // ...
}
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.toImageMask(minThreshold: 0.7)

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

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

let peopleMask = result.toImageMask(minThreshold: 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.toImageMask(of: FritzVisionPeopleClass.person)

Additionally, you can call toImageMask with minThresholdAccepted and threshold arguments as well. This is helpful for dealing with the uncertainty of the model.

  • When minThresholdAccepted is set, any probabilities below that threshold will have an alpha value of 0.
  • When threshold is set, any probabilities above threshold will have an alpha value of 255.
  • Any probabilities between minThresholdAccepted and threshold 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.toImageMask(
    of: FritzVisionPeopleClass.person,
    threshold: 0.7,
    minThresholdAccepted: 0.3)

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.

3. 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
}

4. 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)

5. 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: