How to run tflite model with ARKit session captured image?

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)
}

score 0 · Answer 1 · answered Jul 07 '20 at 16:46

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

        }

}

How to run tflite model with ARKit session captured image?

1 Answers1