我正在尝试让苹果公司在2017年WWDC上展示的Core ML模型样本正常工作。我使用GoogLeNet来尝试对图像进行分类(请参见苹果机器学习页面)。该模型需要一个CVPixelBuffer作为输入。我有一个名为imageSample.jpg的图像用于此演示。我的代码如下:
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)
我总是会在输出中得到意外的运行时错误,而不是图像分类。我转换图像的代码如下:
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 }
我从之前的StackOverflow帖子中获取了这段代码(最后一个回答在这里)。我知道这段代码可能不正确,但我不知道如何自己完成。我认为这是包含错误的部分。模型需要以下类型的输入:Image<RGB,224,224>
回答:
您不需要自己进行大量的图像处理来使用带有图像的Core ML模型——新的Vision框架可以为您完成这些工作。
import Visionimport CoreMLlet 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) }}
WWDC17关于Vision的会议应该会提供更多信息——它将在明天下午举行。