【iOS ARKit】人形提取

为解决人形分离和深度估计问题，ARKit 新增加了 Segmentation Buffer（人体分隔缓冲区）和Estimated Depth Data Buffer（深度估计缓冲区）两个缓冲区。人体分隔缓冲区作用类似于图形渲染管线中的 Stencil Buffer（模板缓冲区），用于区分人形区域与背景区域，它是一个像素级的缓冲区，用于精确地描述人形区域。

人体分隔缓冲区用于标识人形区域，所以可以使用非常简单的结构，如使用1标识该像素是人形区域，而用。标识该像素为背景区。人体分隔缓冲区每帧都更新，所以可以动态地追踪摄像头采集的人形变化。

既然人体分隔缓冲区标识了人形区域，我们也就可以利用该缓冲区提取出场景中的人形以便后续应用，如将人形图像通过网络传输到其他AR设备中，实现类似虚拟会议的效果；或者将人形图像放入虚拟世界中，营造更绚酷的体验；或者对提取的人形图像进行模糊和打马赛克等处理，实现以往只能使用绿幕才能实现的实时人形捕捉效果。

为简单起见，本节我们直接获取人体分隔缓冲区数据并将其保存为图像，关键代码如代码如下所示。

//
//  HumanExtraction.swift
//  ARKitDeamo
//
//  Created by zhaoquan du on 2024/2/4.
//import SwiftUI
import ARKit
import RealityKit
import Combine
import VideoToolbox
import AVFoundationstruct HumanExtraction: View {var viewModel = HumanExtractionViewModel()var arView: ARView {let arView = ARView(frame: .zero)return arView}var body: some View {HumanExtractionContainer(viewModel: viewModel).overlay(VStack{Spacer()Button(action:{viewModel.catchHuman()}) {Text("截取人形").frame(width:120,height:40).font(.body).foregroundColor(.black).background(Color.white).opacity(0.6)}.offset(y:-30).padding(.bottom, 30)}).edgesIgnoringSafeArea(.all)}
}struct HumanExtractionContainer : UIViewRepresentable{var viewModel: HumanExtractionViewModelfunc makeUIView(context: Context) -> some ARView {let arView = ARView(frame: .zero)return arView}func updateUIView(_ uiView: UIViewType, context: Context) {guard ARWorldTrackingConfiguration.supportsFrameSemantics(.personSegmentation) else {return}let config = ARWorldTrackingConfiguration()config.frameSemantics = .personSegmentationuiView.session.delegate = viewModeluiView.session.run(config)}}class HumanExtractionViewModel: NSObject,ARSessionDelegate {var arFrame: ARFrame? = nilfunc session(_ session: ARSession, didUpdate frame: ARFrame) {arFrame = frame}func catchHuman(){if let segmentationBuffer = arFrame?.segmentationBuffer {if let uiImage = UIImage(pixelBuffer: segmentationBuffer)?.rotate(radians: .pi / 2) {UIImageWriteToSavedPhotosAlbum(uiImage, self, #selector(imageSaveHandler(image:didFinishSavingWithError:contextInfo:)), nil)}}}@objc func imageSaveHandler(image:UIImage,didFinishSavingWithError error:NSError?,contextInfo:AnyObject) {if error != nil {print("保存图片出错")} else {print("保存图片成功")}}}extension UIImage {public convenience init?(pixelBuffer:CVPixelBuffer) {var cgimage: CGImage?VTCreateCGImageFromCVPixelBuffer(pixelBuffer, options: nil, imageOut: &cgimage)if let cgimage = cgimage{self.init(cgImage: cgimage)}else{return nil}}func rotate(radians: CGFloat) -> UIImage {let rotatedSize = CGRect(origin: .zero, size: size).applying(CGAffineTransform(rotationAngle: CGFloat(radians))).integral.sizeUIGraphicsBeginImageContext(rotatedSize)if let context = UIGraphicsGetCurrentContext() {let origin = CGPoint(x: rotatedSize.width / 2.0, y: rotatedSize.height / 2.0)context.translateBy(x: origin.x, y: origin.y)context.rotate(by: radians)draw(in: CGRect(x: -origin.y, y: -origin.x, width: size.width, height: size.height))let rotateImage = UIGraphicsGetImageFromCurrentImageContext()UIGraphicsEndImageContext()return rotateImage ?? self}return self}
}

在代码中，人体分隔缓冲区数据每帧都会更新，所以我们需要从 ARFrame 中实时获取值，然后将缓冲区中的数据转换成图像，由于缓冲区中的数据是直接对应硬件摄像头采集的图像数据，为与屏幕显示保持一致，需要对图像进行90°旋转，保存的图像如下右图所示。