Search code examples
swiftaugmented-realityarkitcoremlcvpixelbuffer

Is it possible to change a resolution of captured AR Images in CVPixelBuffer?


I'm using pre-trained CoreML model inside ARKit app. And I'm capturing images from ARCamera and placing them into CVPixelBuffer for processing:

let pixelBuffer: CVPixelBuffer? = sceneView.session.currentFrame?.capturedImage

ARKit can capture pixel buffers in a YCbCr format. To correctly render these images on an iPhone's display, you'll need to access the luma and chroma planes of the pixel buffer and convert full-range YCbCr values to an sRGB using float4x4 ycbcrToRGBTransform matrix. So I understand how to handle a color.

But I'd like to know if I can change a resolution of Captured AR Images in CVPixelBuffer?

How to do it? I need a processing to be as low as possible.


Solution

  • Yes, you can do it. Here is how!

    /**
     Resizes a CVPixelBuffer to a new width and height.
     */
    func resizePixelBuffer(_ pixelBuffer: CVPixelBuffer,
                           width: Int, height: Int) -> CVPixelBuffer? {
        return resizePixelBuffer(pixelBuffer, cropX: 0, cropY: 0,
                                 cropWidth: CVPixelBufferGetWidth(pixelBuffer),
                                 cropHeight: CVPixelBufferGetHeight(pixelBuffer),
                                 scaleWidth: width, scaleHeight: height)
    }
    
    func resizePixelBuffer(_ srcPixelBuffer: CVPixelBuffer,
                           cropX: Int,
                           cropY: Int,
                           cropWidth: Int,
                           cropHeight: Int,
                           scaleWidth: Int,
                           scaleHeight: Int) -> CVPixelBuffer? {
    
        CVPixelBufferLockBaseAddress(srcPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        guard let srcData = CVPixelBufferGetBaseAddress(srcPixelBuffer) else {
            print("Error: could not get pixel buffer base address")
            return nil
        }
        let srcBytesPerRow = CVPixelBufferGetBytesPerRow(srcPixelBuffer)
        let offset = cropY*srcBytesPerRow + cropX*4
        var srcBuffer = vImage_Buffer(data: srcData.advanced(by: offset),
                                      height: vImagePixelCount(cropHeight),
                                      width: vImagePixelCount(cropWidth),
                                      rowBytes: srcBytesPerRow)
    
        let destBytesPerRow = scaleWidth*4
        guard let destData = malloc(scaleHeight*destBytesPerRow) else {
            print("Error: out of memory")
            return nil
        }
        var destBuffer = vImage_Buffer(data: destData,
                                       height: vImagePixelCount(scaleHeight),
                                       width: vImagePixelCount(scaleWidth),
                                       rowBytes: destBytesPerRow)
    
        let error = vImageScale_ARGB8888(&srcBuffer, &destBuffer, nil, vImage_Flags(0))
        CVPixelBufferUnlockBaseAddress(srcPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        if error != kvImageNoError {
            print("Error:", error)
            free(destData)
            return nil
        }
    
        let releaseCallback: CVPixelBufferReleaseBytesCallback = { _, ptr in
            if let ptr = ptr {
                free(UnsafeMutableRawPointer(mutating: ptr))
            }
        }
    
        let pixelFormat = CVPixelBufferGetPixelFormatType(srcPixelBuffer)
        var dstPixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreateWithBytes(nil, scaleWidth, scaleHeight,
                                                  pixelFormat, destData,
                                                  destBytesPerRow, releaseCallback,
                                                  nil, nil, &dstPixelBuffer)
        if status != kCVReturnSuccess {
            print("Error: could not create new pixel buffer")
            free(destData)
            return nil
        }
        return dstPixelBuffer
    }
    

    Usage:

    if let pixelBuffer = sceneView.session.currentFrame?.capturedImage, let resizedBuffer = resizePixelBuffer(pixelBuffer, width: 320, height: 480) {
        //Core Model Processing
    }
    

    Reference: https://github.com/hollance/CoreMLHelpers/tree/master/CoreMLHelpers