How to run tflite model with ARKit session captured image? - ios

Question

I have a tflite model and I want to run model using ARKit session captured image. It's showing source pixel format is invalid. I was able to run tflite model using AVCapture session:
// 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
}
}