Skip to content

Instantly share code, notes, and snippets.

@hradec
Last active December 22, 2023 09:26
Show Gist options
  • Save hradec/6f0dd29e1acfc90ad154588cac1918bd to your computer and use it in GitHub Desktop.
Save hradec/6f0dd29e1acfc90ad154588cac1918bd to your computer and use it in GitHub Desktop.
NSImage extension with functions that convert it to Depth, Disparity, Color and Mask CVPixelBuffer to be used directly in a PhotogrammetrySample object.
// this code is based on the original extension found in this
// other gist: https://gist.github.com/DennisWeidmann/7c4b4bb72062bd1a40c714aa5d95a0d7
// thanks Dennis Weidmann!
extension NSImage {
// function used by all functions below. It shouldn't be used directly
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)
}
}
@hradec
Copy link
Author

hradec commented Jan 12, 2022

To use PhotogrammetrySample, we also need to read the metadata from a HEIC file and the gravity.TXT. So I'm adding the code for those, so you don't have to expend all the time trying to figure those out like I did:

import ImageIO
func readEXIF(file: String) -> [ String: Any ]? {
    var dict:[ String: Any ] = [:]
    if let imageSource = CGImageSourceCreateWithURL(NSURL(fileURLWithPath: file) as CFURL, nil) {
        let imageProperties = CGImageSourceCopyPropertiesAtIndex(imageSource, 0, nil)
        dict = (imageProperties as? [String: Any])!
        let w = dict["PixelWidth"]
        let h = dict["PixelHeight"]
        // I use edited HEIC images that have the camera rotation 
        // applied, so I just exchange w and h here.
        // If you don't need this, just comment it out!
        dict["PixelWidth"] = h
        dict["PixelHeight"] = w
        // I remove the "Depth" info since I'm reading the depth from a TIF image
        dict["Depth"] = nil
        return dict
    }
    return [:]
}

import CoreMotion
func readGravity(file: String) -> CMAcceleration? {
    let gfile = file //file.dropLast(5)+"_gravity.TXT"
    let contents = try! String(contentsOfFile: String(gfile))
    let d = contents.split(separator: ",")
    let gravity = CMAcceleration(x: Double(d[0])!, y: Double(d[1])!, z: Double(d[2])!)
    return gravity
}

I edit my HEIC photos in Natron, wich applies the rotation from the EXIF metadata of the photo automatically. So to use edited HEIC images with the EXIF from the original HEIC photos, I have to exchange the width and height.
If you don't need that, just comment these lines out:

//dict["PixelWidth"] = h
//dict["PixelHeight"] = w

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