Get RGB "CVPixelBuffer" from ARKit

Question

I'm trying to get a CVPixelBuffer in RGB color space from the Apple's ARKit. In func session(_ session: ARSession, didUpdate frame: ARFrame) method of ARSessionDelegate I get an instance of ARFrame. On page Displaying an AR Experience with Metal I found that this pixel buffer is in YCbCr (YUV) color space.

I need to convert this to RGB color space (I actually need CVPixelBuffer and not UIImage). I've found something about color conversion on iOS but I was not able to get this working in Swift 3.

Electroluminescence

What use case do you have to need RGB over YUV? – Castara

I have a custom processing pipeline that requires RGB – Electroluminescence

Ok. If your pipeline involved some kind of opengl / Metal, it's easy to do the conversion there. I'm not sure there's a good way to get RGB directly. – Castara

Answer 1

There's several ways to do this, depending on what you're after. The best way to do this in realtime (to say, render the buffer to a view) is to use a custom shader to convert the YCbCr CVPixelBuffer to RGB.

Using Metal: If you make a new project, select "Augmented Reality App," and select "Metal" for the content technology, the project generated will contain the code and shaders necessary to make this conversion.

Using OpenGL: The GLCameraRipple example from Apple uses an AVCaptureSession to capture the camera, and shows how to map the resulting CVPixelBuffer to GL textures, which are then converted to RGB in shaders (again, provided in the example).

Non Realtime: The answer to this stackoverflow question addresses converting the buffer to a UIImage, and offers a pretty simple way to do it.

Cantlon

Answer 2

I have also stuck on this question for several days. All of the code snippet I could find on the Internet is written in Objective-C rather than Swift, regarding converting CVPixelBuffer to UIImage.

Finally, the following code snippet works perfect for me, to convert a YUV image to either JPG or PNG file format, and then you can write it to the local file in your application.

func pixelBufferToUIImage(pixelBuffer: CVPixelBuffer) -> UIImage {
    let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
    let context = CIContext(options: nil)
    let cgImage = context.createCGImage(ciImage, from: ciImage.extent)
    let uiImage = UIImage(cgImage: cgImage!)
    return uiImage
}

Garrick

This really does not answer the question. – Electroluminescence

This is the best answer as it properly handles colorspace and gamma conversion, but keep in mind that allocating a CIContext like that on every call with be REALLY REALLY slow. Hold a ref to the CIContext for the lifetime of repeated calls and your app will perform effectively. If you are interested in a full Metal impl of RGB->YCbCr conversion then have a look at: github.com/mdejong/MetalBT709Decoder – Durra

ARKit with SceneKit provides a snapshot method to get the current frame to UIImage directly. – Quebec

Ditto to what @Durra said. I tried doing this on captured video frames and my app would crash almost immediately, possibly because it was running out of memory. I moved the CIContext out of my function so I could reuse it, and now this approach works great. – Yaupon

Answer 3

You may want the Accelerate framework's image conversion functions. Perhaps a combination of vImageConvert_420Yp8_Cb8_Cr8ToARGB8888 and vImageConvert_ARGB8888toRGB888 (If you don't want the alpha channel). In my experience these work in real time.

Dewdrop

Answer 4

The docs explicitly says that you need to access the luma and chroma planes:

ARKit captures pixel buffers in a planar YCbCr format (also known as YUV) format. To render these images on a device display, you'll need to access the luma and chroma planes of the pixel buffer and convert pixel values to an RGB format.

So there's no way to directly get the RGB planes and you'll have to handle this in your shaders, either in Metal or openGL as described by @joshue

Castara

Answer 5

Struggled a long while with this as well and I've ended up writing the following code, which works for me:

// Helper macro to ensure pixel values are bounded between 0 and 255
#define clamp(a) (a > 255 ? 255 : (a < 0 ? 0 : a));

- (void)processImageBuffer:(CVImageBufferRef)imageBuffer
{
    OSType type  = CVPixelBufferGetPixelFormatType(imageBuffer);
    if (type == kCVPixelFormatType_420YpCbCr8BiPlanarFullRange)
    {
        CVPixelBufferLockBaseAddress(imageBuffer, 0);
        // We know the return format of the base address based on the YpCbCr8BiPlanarFullRange format (as per doc)
        StandardBuffer baseAddress = (StandardBuffer)CVPixelBufferGetBaseAddress(imageBuffer);

        // Get the number of bytes per row for the pixel buffer, width and height
        size_t bytesPerRow = CVPixelBufferGetBytesPerRow(imageBuffer);
        size_t width = CVPixelBufferGetWidth(imageBuffer);
        size_t height = CVPixelBufferGetHeight(imageBuffer);

        // Get buffer info and planar pixel data
        CVPlanarPixelBufferInfo_YCbCrBiPlanar *bufferInfo = (CVPlanarPixelBufferInfo_YCbCrBiPlanar *)baseAddress;
        uint8_t* cbrBuff = (uint8_t *)CVPixelBufferGetBaseAddressOfPlane(imageBuffer, 1);
        // This just moved the pointer past the offset
        baseAddress = (uint8_t *)CVPixelBufferGetBaseAddressOfPlane(imageBuffer, 0);
        int bytesPerPixel = 4;
        uint8_t *rgbData =  rgbFromYCrCbBiPlanarFullRangeBuffer(baseAddress,
                                                                cbrBuff,
                                                                bufferInfo,
                                                                width,
                                                                height,
                                                                bytesPerRow);

        [self doStuffOnRGBBuffer:rgbData width:width height:height bitsPerComponent:8 bytesPerPixel:bytesPerPixel bytesPerRow:bytesPerRow];

        free(rgbData);
        CVPixelBufferUnlockBaseAddress(imageBuffer, 0);
    }
    else
    {
        NSLog(@"Unsupported image buffer type");
    }
}

uint8_t * rgbFromYCrCbBiPlanarFullRangeBuffer(uint8_t *inBaseAddress,
                                              uint8_t *cbCrBuffer,
                                              CVPlanarPixelBufferInfo_YCbCrBiPlanar * inBufferInfo,
                                              size_t inputBufferWidth,
                                              size_t inputBufferHeight,
                                              size_t inputBufferBytesPerRow)
{
    int bytesPerPixel = 4;
    NSUInteger yPitch = EndianU32_BtoN(inBufferInfo->componentInfoY.rowBytes);
    uint8_t *rgbBuffer = (uint8_t *)malloc(inputBufferWidth * inputBufferHeight * bytesPerPixel);
    NSUInteger cbCrPitch = EndianU32_BtoN(inBufferInfo->componentInfoCbCr.rowBytes);
    uint8_t *yBuffer = (uint8_t *)inBaseAddress;

    for(int y = 0; y < inputBufferHeight; y++)
    {
        uint8_t *rgbBufferLine = &rgbBuffer[y * inputBufferWidth * bytesPerPixel];
        uint8_t *yBufferLine = &yBuffer[y * yPitch];
        uint8_t *cbCrBufferLine = &cbCrBuffer[(y >> 1) * cbCrPitch];
        for(int x = 0; x < inputBufferWidth; x++)
        {
            int16_t y = yBufferLine[x];
            int16_t cb = cbCrBufferLine[x & ~1] - 128;
            int16_t cr = cbCrBufferLine[x | 1] - 128;

            uint8_t *rgbOutput = &rgbBufferLine[x*bytesPerPixel];

            int16_t r = (int16_t)roundf( y + cr *  1.4 );
            int16_t g = (int16_t)roundf( y + cb * -0.343 + cr * -0.711 );
            int16_t b = (int16_t)roundf( y + cb *  1.765);

            // ABGR image representation
            rgbOutput[0] = 0Xff;
            rgbOutput[1] = clamp(b);
            rgbOutput[2] = clamp(g);
            rgbOutput[3] = clamp(r);
        }
    }

    return rgbBuffer;
}

Korella

Is there a way to leverage this code in a SWIFT application? – Demotic