如何使用ARKit会话捕获的图像运行tflite模型?

Question

我有一个tflite模型，我想使用 ARKit 会话捕获的图像运行模型。它显示源像素格式无效。我能够使用 AVCapture 会话运行tflite模型：

// Not working with ARKit frame
func session(_ session: ARSession, didUpdate frame: ARFrame) {
guard currentBuffer == nil, case .normal = frame.camera.trackingState else { return }
runModel(with : currentBuffer)
}
// Working fine with AVCapture session
func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer)
runModel(with : pixelBuffer)
}

Answer 1

let context = CIContext()

func session(_ session: ARSession, didUpdate frame: ARFrame) {
guard currentBuffer == nil, case .normal = frame.camera.trackingState else {
return
}

guard let scaledPixelBuffer = CIImage(cvPixelBuffer: frame.capturedImage)
.oriented(.right)
.resize(size: CGSize(width: 256, height: 256))
.toPixelBuffer(context: context)

else {
return
}
runModel(with : scaledPixelBuffer)
}



extension CIImage {

func resize(size : CGSize) -> CIImage {
let scale = min(size.width, size.height) / min(self.extent.size.width, self.extent.size.height)
let resizedImage = self.transformed(by: CGAffineTransform(scaleX: scale, y: scale))
let width = resizedImage.extent.size.width
let height = resizedImage.extent.size.height
let xOffset = (CGFloat(width) - size.width) / 2.0
let yOffset = (CGFloat(height) - size.height) / 2.0
let rect = CGRect(x: xOffset, y: yOffset, width: size.width, height: size.height)
return resizedImage
.clamped(to: rect)
.cropped(to: CGRect(x: 0, y: 0, width: size.width, height: size.height))

}

func toPixelBuffer(context : CIContext, size inSize:CGSize? = nil, gray : Bool = true) -> CVPixelBuffer? {
let attributes = [
kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue ] as CFDictionary
var nullablePixelBuffer : CVPixelBuffer? = nil
let status = CVPixelBufferCreate(kCFAllocatorDefault, Int(self.extent.size.width), Int(self.extent.size.height), kCVPixelFormatType_32BGRA, attributes, &nullablePixelBuffer)
guard status == kCVReturnSuccess, let pixelBuffer = nullablePixelBuffer else {
return nil
}
CVPixelBufferLockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
context.render(self, to: pixelBuffer, bounds: CGRect(x: 0, y: 0, width: self.extent.size.width, height: self.extent.size.height), colorSpace: gray ? CGColorSpaceCreateDeviceGray() : self.colorSpace)
CVPixelBufferUnlockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
return pixelBuffer
}
}