【iOS ARKit】人形提取

news2024/10/6 16:31:14

为解决人形分离和深度估计问题，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 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
    }
}

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