Reimplementing iOS Photos app background removal / image segmentation in Swift?

Question

I have searched all over for examples of how to mimic the newer iOS Photos app capability to lift an object from it's background in my own SwiftUI project, but haven't found anything better than this examples:

https://betterprogramming.pub/coreml-image-segmentation-background-remove-ca11e6f6a083

which does a decent job, but the results don't compare to the quality of segmentation pulled off from within the Photos app (starting with iOS 16, I believe). Presumably all the functionality is available in XCode via the various Core ML models, but I can't find any examples of how to pull of similarly clean image segmentation.

When implementing the code found in this example (https://betterprogramming.pub/coreml-image-segmentation-background-remove-ca11e6f6a083 ), which uses the DeeplabV3 image segmentation functions, i get the following results:

DeeplabV3 segmentation

Here is what the Photos App is able to accomplish for comparison:

Photos App segmentation

I'm very new to Swift programming, so maybe I'm being a bit naive regarding accessibility to default iOS app methods. Hoping someone can point me in the right direction.

Answer 1

The function you want is VNGenerateForegroundInstanceMaskRequest. Here's an example wrapper that removes the background from an image:

import Foundation
import CoreImage
import Vision

public struct BackgroundRemover {
    enum ImageExtractError: Error {
        case noResultFromForegroundMaskRequest
    }
        
    public static func removeBackground(from image: CIImage) throws -> CIImage {
        let request = VNGenerateForegroundInstanceMaskRequest()

        let handler = VNImageRequestHandler(ciImage: image, options: [:])
        try handler.perform([request])

        guard let result = request.results?.first else {
            throw ImageExtractError.noResultFromForegroundMaskRequest
        }

        let maskedImage = try result.generateMaskedImage(
            ofInstances: result.allInstances,
            from: handler,
            croppedToInstancesExtent: true
        )
        
        let image = CIImage(cvPixelBuffer: maskedImage)
        return image
    }
}

This works with CIImage so it will work on both iOS and macOS (you can convert to/from UIImage using image.cgImage or UIImage(cgImage: someImage)).