iosswiftcore-imagecifilterimagefilter

iOS Core Filter CIPersonSegmentation creates a mask whose position is different from original. How to correctly position it?

I have the following image:

enter image description here

I am using the below code to create a red mask from CIPersonSegmentation and then change the background color to cyan:

WARNING: CIPersonSegmentation requires physical device. It uses Vision API which I believe is not supported on simulator.

if let image = UIImage(named: "demo9"), let editted = applyPersonFilter(to: image) {
    imageView.image = editted
}

func applyPersonFilter(to image: UIImage) -> UIImage? {
    guard let ciImage = CIImage(image: image) else { return nil }
    
    print("ciImage: \(ciImage.extent)") // prints (0.0, 0.0, 2882.0, 1300.0)
    
    let context = CIContext(options: nil)
    
    let personMaskFilter = CIFilter(name: "CIPersonSegmentation", parameters: [
        kCIInputImageKey: ciImage,
        "inputQualityLevel": 0
    ])
    
    let maskImage = personMaskFilter?.outputImage?.cropped(to: ciImage.extent) as? CIImage
    
    print("maskImage: \(maskImage?.extent)") // prints (0.0, 0.0, 2016.0, 1300.0)
    
    let solidColorFilter = CIFilter(name: "CIConstantColorGenerator")
    solidColorFilter?.setValue(CIColor.cyan, forKey: kCIInputColorKey)

    let blendFilter = CIFilter(name: "CIBlendWithRedMask", parameters: [kCIInputImageKey: ciImage, kCIInputBackgroundImageKey: solidColorFilter?.outputImage?.cropped(to: ciImage.extent) as Any, kCIInputMaskImageKey: maskImage as Any])
        
    guard let outputImage = blendFilter?.outputImage else {
        return nil
    }
    
    print("blendFilter: \(outputImage.extent)") // prints (0.0, 0.0, 2882.0, 1300.0)
    
    guard let cgImage = context.createCGImage(outputImage, from: outputImage.extent) else { return nil }
    
    return UIImage(cgImage: cgImage)
}

The output it creates is this:

enter image description here

As you can see, the mask position is not correct. This is because the red mask created by CIPersonSegmentation has a different extent than the original image.

How to correctly apply this taking into account the different extend of the mask?

Solution

  • The CIPersonSegmentation filter is ML-based under the hood. And those ML-models have a maximum input size. My guess is that 2016 is the maximum size for that quality level.

    The documentation of qualityLevel says

    The size and quality of the resulting segmentation mask. 0 is accurate, 1 is balanced, and 2 is fast.

    So maybe, if you increase the qualityLevel, the generated mask will also be bigger.

    However, to be on the safe side, you can simply stretch the mask to the extent of the original image using a CGAffineTransform:

    let scaleX = ciImage.extent.width / maskImage.extent.width
let scaleY = ciImage.extent.height / maskImage.extent.height
let transform = CGAffineTransform(scaleX: scaleX, y: scaleY)
let scaledMask = maskImage.transformed(by: transform)