iOS MachineLearning 系列（9）—— 人物蒙版图生成

作者：珲少

2023-04-27
上海
本文字数：1775 字
阅读完需：约 6 分钟

人物蒙版图能力是 Vision 框架在 iOS 15 中新增的功能，这个功能可以将图片中的人物按照轮廓生成无光蒙版。无光蒙版在实际业务中非常有用，使用此蒙版可以方便的将人物从图片中提取出来，然后和其他的背景图进行合成。

1 - 人物蒙版的提取

首先，人物蒙版的提取非常简单，使用 VNGeneratePersonSegmentationRequest 创建蒙版分析请求，如下：

private lazy var personRequest: VNGeneratePersonSegmentationRequest = {    let request = VNGeneratePersonSegmentationRequest { request, error in        DispatchQueue.main.async {            self.drawTask(request: request as! VNGeneratePersonSegmentationRequest)        }    }    return request}()

复制代码

使用 VNImageRequestHandler 来触发请求即可：

// 要分析的图片资源let image = UIImage(named: "image7")!
// 图像分析请求lazy var imageRequestHandler = VNImageRequestHandler(ciImage: CIImage(cgImage: image.cgImage!),                                                orientation: .up,                                                options: [:])

复制代码

请求的结果 VNPixelBufferObservation 中会封装蒙版图片 CVPixelBuffer 数据，如下：

open class VNPixelBufferObservation : VNObservation {    // 分析出的蒙版数据    open var pixelBuffer: CVPixelBuffer { get }    // 分析所使用的模型    open var featureName: String? { get }}

复制代码

处理请求结果如下：

private func drawTask(request: VNGeneratePersonSegmentationRequest) {
    for result in request.results ?? [] {        // 创建CIImage实例        let ciImage = CIImage(cvPixelBuffer: result.pixelBuffer)        // 默认返回的蒙版为黑白两色，这里讲所有黑色替换成透明色        let filter = CIFilter(name: "CIMaskToAlpha", parameters: [kCIInputImageKey: ciImage])        guard let outputImage = filter?.outputImage else { return }        let context = CIContext(options: nil)        let cgImage = context.createCGImage(outputImage, from: outputImage.extent)        let uiImage = UIImage(cgImage: cgImage!)                let v = UIImageView(image: uiImage)        v.frame = imageView.frame        v.backgroundColor = .black        v.frame.origin.y += imageView.frame.height        v.image = uiImage        view.addSubview(v)    }}

复制代码

效果如下图：

2 - 进行人物提取

获取到了蒙版，提取人物将很简单，修改上述代码如下：

private func drawTask(request: VNGeneratePersonSegmentationRequest) {
    for result in request.results ?? [] {        print(result.pixelBuffer)        let ciImage = CIImage(cvPixelBuffer: result.pixelBuffer)                let filter = CIFilter(name: "CIMaskToAlpha", parameters: [kCIInputImageKey: ciImage])                guard let outputImage = filter?.outputImage else { return }        let context = CIContext(options: nil)        let cgImage = context.createCGImage(outputImage, from: outputImage.extent)        let uiImage = UIImage(cgImage: cgImage!)                        let v = UIImageView(image: uiImage)        v.frame = imageView.frame        v.backgroundColor = .blue        v.frame.origin.y += imageView.frame.height
        v.image = uiImage        view.addSubview(v)                                let v2 = UIImageView(image: image)        v2.frame = v.frame        v2.frame.origin.y += v.frame.height                let mask =  UIImageView(image: uiImage)        mask.frame = CGRect(x: 0, y: 0, width: v2.frame.size.width, height: v2.frame.size.height)        mask.backgroundColor = .clear        // 使用蒙版截取        v2.mask = mask        view.addSubview(v2)    }}