Convert NSImage to CVPixelBuffer

NSImageExtension.swift

extension NSImage {
    func pixelBuffer() -> CVPixelBuffer? {
        let width = self.size.width
        let height = self.size.height
        let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
                     kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                         Int(width),
                                         Int(height),
                                         kCVPixelFormatType_32ARGB,
                                         attrs,
                                         &pixelBuffer)
        
        guard let resultPixelBuffer = pixelBuffer, status == kCVReturnSuccess else {
            return nil
        }
        
        CVPixelBufferLockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        let pixelData = CVPixelBufferGetBaseAddress(resultPixelBuffer)
        
        let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
        guard let context = CGContext(data: pixelData,
                                      width: Int(width),
                                      height: Int(height),
                                      bitsPerComponent: 8,
                                      bytesPerRow: CVPixelBufferGetBytesPerRow(resultPixelBuffer),
                                      space: rgbColorSpace,
                                      bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue) else {return nil}
        
        context.translateBy(x: 0, y: height)
        context.scaleBy(x: 1.0, y: -1.0)
        
        let graphicsContext = NSGraphicsContext(cgContext: context, flipped: false)
        NSGraphicsContext.saveGraphicsState()
        NSGraphicsContext.current = graphicsContext
        draw(in: CGRect(x: 0, y: 0, width: width, height: height))
        NSGraphicsContext.restoreGraphicsState()
        
        CVPixelBufferUnlockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        
        return resultPixelBuffer
    }
}

In case someone else needs to convert NSImage to proper depthData, disparityData ,color and mask CVPixelBuffers that can be passed directly to a PhotogrammetrySample in ObjectCapture, I've modified the original code slightly to add the these functions:

extension NSImage {
    // return the NSImage as a color 32bit Color CVPixelBuffer
    // function used by depthPixelBuffer and disparityPixelBuffer to actually crate the CVPixelBuffer
    func __toPixelBuffer(PixelFormatType: OSType) -> CVPixelBuffer? {
        var bitsPerC   = 8
        var colorSpace = CGColorSpaceCreateDeviceRGB()
        var bitmapInfo = CGImageAlphaInfo.noneSkipFirst.rawValue
        
        // if we need depth/disparity
        if PixelFormatType == kCVPixelFormatType_DepthFloat32 || PixelFormatType == kCVPixelFormatType_DisparityFloat32 {
            bitsPerC   = 32
            colorSpace = CGColorSpaceCreateDeviceGray()
            bitmapInfo = CGImageAlphaInfo.none.rawValue | CGBitmapInfo.floatComponents.rawValue
        }
        // if we need mask
        else if PixelFormatType == kCVPixelFormatType_OneComponent8 {
            bitsPerC   = 8
            colorSpace = CGColorSpaceCreateDeviceGray()
            bitmapInfo = CGImageAlphaInfo.alphaOnly.rawValue
        }
        
        let width  = Int(self.size.width)
        let height = Int(self.size.height)
        let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
                     kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary

        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault, width, height, PixelFormatType, attrs, &pixelBuffer)
        guard let resultPixelBuffer = pixelBuffer, status == kCVReturnSuccess else {
            return nil
        }
        
        CVPixelBufferLockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        guard let context = CGContext(data: CVPixelBufferGetBaseAddress(resultPixelBuffer),
                                      width: width,
                                      height: height,
                                      bitsPerComponent: bitsPerC,
                                      bytesPerRow: CVPixelBufferGetBytesPerRow(resultPixelBuffer),
                                      space: colorSpace,
                                      bitmapInfo: bitmapInfo)
        else {
            return nil
        }
        
        // context.translateBy(x: 0, y: height)
        // context.scaleBy(x: 1.0, y: -1.0)
        
        let graphicsContext = NSGraphicsContext(cgContext: context, flipped: false)
        NSGraphicsContext.saveGraphicsState()
        NSGraphicsContext.current = graphicsContext
        draw(in: CGRect(x: 0, y: 0, width: width, height: height))
        NSGraphicsContext.restoreGraphicsState()
        
        CVPixelBufferUnlockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
        
        return resultPixelBuffer
    }

    // return the NSImage as a color 32bit Color CVPixelBuffer
    func colorPixelBuffer() -> CVPixelBuffer? {
        return __toPixelBuffer(PixelFormatType: kCVPixelFormatType_32ARGB)
    }

    func maskPixelBuffer() -> CVPixelBuffer? {
        return __toPixelBuffer(PixelFormatType: kCVPixelFormatType_OneComponent8)
    }

    // return NSImage as a 32bit depthData CVPixelBuffer
    func depthPixelBuffer() -> CVPixelBuffer? {
        return __toPixelBuffer(PixelFormatType: kCVPixelFormatType_DepthFloat32)
    }

    // return NSImage as a 32bit disparityData CVPixelBuffer
    func disparityPixelBuffer() -> CVPixelBuffer? {
        return __toPixelBuffer(PixelFormatType: kCVPixelFormatType_DisparityFloat32)
    }
}

It can also be found in this gist: https://gist.github.com/hradec/6f0dd29e1acfc90ad154588cac1918bd