Skip to content

Instantly share code, notes, and snippets.

@DennisWeidmann
Last active January 9, 2023 10:10
Show Gist options
  • Save DennisWeidmann/7c4b4bb72062bd1a40c714aa5d95a0d7 to your computer and use it in GitHub Desktop.
Save DennisWeidmann/7c4b4bb72062bd1a40c714aa5d95a0d7 to your computer and use it in GitHub Desktop.
Convert NSImage to CVPixelBuffer
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
}
}
@JamesHowlettLA
Copy link

the followed line should be removed:

context.translateBy(x: 0, y: height)
context.scaleBy(x: 1.0, y: -1.0)

@hradec
Copy link

hradec commented Jan 8, 2022

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment