【iOS ARKit】人形提取

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

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

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

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

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

//
//  HumanExtraction.swift
//  ARKitDeamo
//
//  Created by zhaoquan du on 2024/2/4.
//
 
import SwiftUI
import ARKit
import RealityKit
import Combine
import VideoToolbox
import AVFoundation
 
struct 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: HumanExtractionViewModel
     
     
    func 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 = .personSegmentation
        uiView.session.delegate = viewModel
        uiView.session.run(config)
    }
     
     
     
}
 
class HumanExtractionViewModel: NSObject,ARSessionDelegate {
    var arFrame: ARFrame? = nil
    func 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.size
        UIGraphicsBeginImageContext(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
    }
}