26

Apple's new CoreML framework has a prediction function that takes a CVPixelBuffer. In order to classify a UIImage a conversion must be made between the two.

Conversion code I got from an Apple Engineer:

1  // image has been defined earlier
2
3     var pixelbuffer: CVPixelBuffer? = nil
4   
5     CVPixelBufferCreate(kCFAllocatorDefault, Int(image.size.width), Int(image.size.height), kCVPixelFormatType_OneComponent8, nil, &pixelbuffer)
6     CVPixelBufferLockBaseAddress(pixelbuffer!, CVPixelBufferLockFlags(rawValue:0))
7   
8     let colorspace = CGColorSpaceCreateDeviceGray()
9     let bitmapContext = CGContext(data: CVPixelBufferGetBaseAddress(pixelbuffer!), width: Int(image.size.width), height: Int(image.size.height), bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelbuffer!), space: colorspace, bitmapInfo: 0)!
10  
11    bitmapContext.draw(image.cgImage!, in: CGRect(x: 0, y: 0, width: image.size.width, height: image.size.height))

This solution is in swift and is for a grayscale image. Changes that must be made depending on the type of image are:

  • Line 5 | kCVPixelFormatType_OneComponent8 to another OSType (kCVPixelFormatType_32ARGB for RGB)
  • Line 8 | colorSpace to another CGColorSpace (CGColorSpaceCreateDeviceRGB for RGB)
  • Line 9 | bitsPerComponent to the number of bits per pixel of memory (32 for RGB)
  • Line 9 | bitmapInfo to a nonzero CGBitmapInfo property (kCGBitmapByteOrderDefault is the default)
Ryan
  • 269
  • 1
  • 3
  • 3
  • I think `CVPixelBuffer` is just a convention for images from camera, there are also `ciImage`, `cgImage` available – Willjay Jun 09 '17 at 16:51
  • Why not use the Vision Framework? You can (1) initialize a `VNCoreMLModel` object, then a (2) `VNCoreMLRequest` with a completion handler, then use a (3) `VNImageRequestHandler` - which can take a `CIImage` or a file URL. The results are a collection of `VNClassificationObservation` objects. Now of course, if you want to convert a `UIImage` to a `CVPixelBuffer` - which you really **didn't** say anywhere **inside** your question that didn't have a question :-) - you can simply search the internet for that and find this: https://gist.github.com/cieslak/743f9321834c5a40597afa1634a48343 –  Jun 09 '17 at 17:40
  • @dfd Yeah so I did want to convert a `UIImage` to a `CVPixelBuffer` for the purposes of using a CoreML model, but I kindly had this problem solved by an Apple engineer at WWDC with the above code. Given that several other people at the conference have had the same problem, I figured I'd share the solution which achieves its purpose with much more simplicity than that github solution. Thanks for the suggestions though! – Ryan Jun 09 '17 at 18:00
  • 1
    @Ryan: Are you asking a question or do you want to share a solution? In the latter case you should post it as a self-answered question: https://stackoverflow.com/help/self-answer. – Martin R Jun 09 '17 at 18:18
  • +1 to Martin R, please move the second half of your question into an answer to the question, allowing other people to answer too! – Alex Brown Jun 10 '17 at 09:22

1 Answers1

48

You can take a look at this tutorial https://www.hackingwithswift.com/whats-new-in-ios-11, code is in Swift 4

func buffer(from image: UIImage) -> CVPixelBuffer? {
  let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
  var pixelBuffer : CVPixelBuffer?
  let status = CVPixelBufferCreate(kCFAllocatorDefault, Int(image.size.width), Int(image.size.height), kCVPixelFormatType_32ARGB, attrs, &pixelBuffer)
  guard (status == kCVReturnSuccess) else {
    return nil
  }

  CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
  let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)

  let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
  let context = CGContext(data: pixelData, width: Int(image.size.width), height: Int(image.size.height), bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)

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

  UIGraphicsPushContext(context!)
  image.draw(in: CGRect(x: 0, y: 0, width: image.size.width, height: image.size.height))
  UIGraphicsPopContext()
  CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))

  return pixelBuffer
}
onmyway133
  • 45,645
  • 31
  • 257
  • 263