Convertir imagen en CVPixelBuffer para Machine Learning Swift

Estoy tratando de obtener los modelos Core ML de muestra de Apple que se demostraron en la WWDC 2017 para funcionar correctamente. Estoy usando GoogLeNet para tratar de clasificar imágenes (consulte la página de Apple Machine Learning ). El modelo toma un CVPixelBuffer como entrada. Tengo una imagen llamada imageSample.jpg que estoy usando para esta demostración. Mi código está abajo:

var sample = UIImage(named: "imageSample")?.cgImage let bufferThree = getCVPixelBuffer(sample!) let model = GoogLeNetPlaces() guard let output = try? model.prediction(input: GoogLeNetPlacesInput.init(sceneImage: bufferThree!)) else { fatalError("Unexpected runtime error.") } print(output.sceneLabel) 

Siempre obtengo el error de tiempo de ejecución inesperado en la salida en lugar de una clasificación de imagen. Mi código para convertir la imagen está a continuación:

 func getCVPixelBuffer(_ image: CGImage) -> CVPixelBuffer? { let imageWidth = Int(image.width) let imageHeight = Int(image.height) let attributes : [NSObject:AnyObject] = [ kCVPixelBufferCGImageCompatibilityKey : true as AnyObject, kCVPixelBufferCGBitmapContextCompatibilityKey : true as AnyObject ] var pxbuffer: CVPixelBuffer? = nil CVPixelBufferCreate(kCFAllocatorDefault, imageWidth, imageHeight, kCVPixelFormatType_32ARGB, attributes as CFDictionary?, &pxbuffer) if let _pxbuffer = pxbuffer { let flags = CVPixelBufferLockFlags(rawValue: 0) CVPixelBufferLockBaseAddress(_pxbuffer, flags) let pxdata = CVPixelBufferGetBaseAddress(_pxbuffer) let rgbColorSpace = CGColorSpaceCreateDeviceRGB(); let context = CGContext(data: pxdata, width: imageWidth, height: imageHeight, bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(_pxbuffer), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.premultipliedFirst.rawValue) if let _context = context { _context.draw(image, in: CGRect.init(x: 0, y: 0, width: imageWidth, height: imageHeight)) } else { CVPixelBufferUnlockBaseAddress(_pxbuffer, flags); return nil } CVPixelBufferUnlockBaseAddress(_pxbuffer, flags); return _pxbuffer; } return nil } 

Obtuve este código de una publicación anterior de StackOverflow (última respuesta aquí ). Reconozco que el código puede no ser correcto, pero no tengo idea de cómo hacerlo yo mismo. Creo que esta es la sección que contiene el error. El modelo requiere el siguiente tipo de entrada: Image

No es necesario que se haga un montón de imágenes para usar un modelo Core ML con una imagen: el nuevo marco Vision lo puede hacer por usted.

 import Vision import CoreML let model = try VNCoreMLModel(for: MyCoreMLGeneratedModelClass().model) let request = VNCoreMLRequest(model: model, completionHandler: myResultsMethod) let handler = VNImageRequestHandler(url: myImageURL) handler.perform([request]) func myResultsMethod(request: VNRequest, error: Error?) { guard let results = request.results as? [VNClassificationObservation] else { fatalError("huh") } for classification in results { print(classification.identifier, // the scene label classification.confidence) } } 

La sesión de WWDC17 sobre Vision debería tener un poco más de información, es mañana por la tarde.

Puede usar un CoreML puro, pero debe cambiar el tamaño de una imagen a (224,224)

  DispatchQueue.global(qos: .userInitiated).async { // Resnet50 expects an image 224 x 224, so we should resize and crop the source image let inputImageSize: CGFloat = 224.0 let minLen = min(image.size.width, image.size.height) let resizedImage = image.resize(to: CGSize(width: inputImageSize * image.size.width / minLen, height: inputImageSize * image.size.height / minLen)) let cropedToSquareImage = resizedImage.cropToSquare() guard let pixelBuffer = cropedToSquareImage?.pixelBuffer() else { fatalError() } guard let classifierOutput = try? self.classifier.prediction(image: pixelBuffer) else { fatalError() } DispatchQueue.main.async { self.title = classifierOutput.classLabel } } // ... extension UIImage { func resize(to newSize: CGSize) -> UIImage { UIGraphicsBeginImageContextWithOptions(CGSize(width: newSize.width, height: newSize.height), true, 1.0) self.draw(in: CGRect(x: 0, y: 0, width: newSize.width, height: newSize.height)) let resizedImage = UIGraphicsGetImageFromCurrentImageContext()! UIGraphicsEndImageContext() return resizedImage } func cropToSquare() -> UIImage? { guard let cgImage = self.cgImage else { return nil } var imageHeight = self.size.height var imageWidth = self.size.width if imageHeight > imageWidth { imageHeight = imageWidth } else { imageWidth = imageHeight } let size = CGSize(width: imageWidth, height: imageHeight) let x = ((CGFloat(cgImage.width) - size.width) / 2).rounded() let y = ((CGFloat(cgImage.height) - size.height) / 2).rounded() let cropRect = CGRect(x: x, y: y, width: size.height, height: size.width) if let croppedCgImage = cgImage.cropping(to: cropRect) { return UIImage(cgImage: croppedCgImage, scale: 0, orientation: self.imageOrientation) } return nil } 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) UIGraphicsPushContext(context) self.draw(in: CGRect(x: 0, y: 0, width: width, height: height)) UIGraphicsPopContext() CVPixelBufferUnlockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0)) return resultPixelBuffer } } 

El tamaño de imagen esperado para las entradas que puede encontrar en el archivo mimodel : enter image description here

Un proyecto de demostración que usa variantes de CoreML y Vision puras que puedes encontrar aquí: https://github.com/handsomecode/iOS11-Demos/tree/coreml_vision/CoreML/CoreMLDemo