将图像转换为 CVPixelBuffer 用于机器学习 Swift [英] Convert Image to CVPixelBuffer for Machine Learning Swift

问题描述

我试图让在 2017 年 WWDC 上演示的 Apple 示例 Core ML 模型正常运行.我正在使用 GoogLeNet 尝试对图像进行分类(请参阅 Apple 机器学习页面).该模型将 CVPixelBuffer 作为输入.我有一个名为 imageSample.jpg 的图像，我在这个演示中使用了它.我的代码如下:

 var sample = UIImage(named: "imageSample")?.cgImage让 bufferThree = getCVPixelBuffer(sample!)让模型 = GoogLeNetPlaces()守卫让输出=尝试?model.prediction(input: GoogLeNetPlacesInput.init(sceneImage: bufferThree!)) else {致命错误(意外的运行时错误.")}打印(输出.场景标签)

我总是在输出中遇到意外的运行时错误，而不是图像分类.我转换图像的代码如下:

func getCVPixelBuffer(_ image: CGImage) ->CVPixelBuffer?{让 imageWidth = Int(image.width)让 imageHeight = Int(image.height)让属性:[NSObject:AnyObject] = [kCVPixelBufferCGImageCompatibilityKey : 作为 AnyObject 为真，kCVPixelBufferCGBitmapContextCompatibilityKey : true as AnyObject]var pxbuffer:CVPixelBuffer?= 零CVPixelBufferCreate(kCFAllocatorDefault,图像宽度，图像高度，kCVPixelFormatType_32ARGB,作为 CFDictionary 的属性?，&pxbuffer)如果让 _pxbuffer = pxbuffer {让标志 = CVPixelBufferLockFlags(rawValue: 0)CVPixelBufferLockBaseAddress(_pxbuffer, flags)让 pxdata = CVPixelBufferGetBaseAddress(_pxbuffer)让 rgbColorSpace = CGColorSpaceCreateDeviceRGB();让上下文 = CGContext(data: pxdata,宽度:图像宽度，高度:图像高度，bitsPerComponent: 8,bytesPerRow: CVPixelBufferGetBytesPerRow(_pxbuffer),空间:rgbColorSpace，位图信息:CGImageAlphaInfo.premultipliedFirst.rawValue)如果让 _context = 上下文 {_context.draw(image, in: CGRect.init(x: 0, y: 0, width: imageWidth, height: imageHeight))}别的 {CVPixelBufferUnlockBaseAddress(_pxbuffer, flags);返回零}CVPixelBufferUnlockBaseAddress(_pxbuffer, flags);返回_pxbuffer；}返回零}

我从之前的 StackOverflow 帖子中得到了这个代码(最后一个答案 here).我认识到代码可能不正确，但我不知道如何自己做到这一点.我相信这是包含错误的部分.该模型需要以下类型的输入:Image

解决方案

你不需要自己做一堆图像处理来使用带有图像的 Core ML 模型——新的 Vision 框架 可以为您做到这一点.

导入视觉导入 CoreML让模型 = 尝试 VNCoreMLModel(for: MyCoreMLGeneratedModelClass().model)让请求 = VNCoreMLRequest(模型:模型，completionHandler:myResultsMethod)让处理程序 = VNImageRequestHandler(url: myImageURL)handler.perform([请求])func myResultsMethod(请求:VNRequest，错误:错误?){守卫让结果= request.result as?【VN分类观察】其他 { 致命错误(嗯")}用于结果分类{print(classification.identifier,//场景标签分类.置信度)}}

WWDC17 Vision 会议应该有更多信息 - 明天下午.>

I am trying to get Apple's sample Core ML Models that were demoed at the 2017 WWDC to function correctly. I am using the GoogLeNet to try and classify images (see the Apple Machine Learning Page). The model takes a CVPixelBuffer as an input. I have an image called imageSample.jpg that I'm using for this demo. My code is below:

        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)

I am always getting the unexpected runtime error in the output rather than an image classification. My code to convert the image is below:

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
    }

I got this code from a previous StackOverflow post (last answer here). I recognize that the code may not be correct, but I have no idea of how to do this myself. I believe that this is the section that contains the error. The model calls for the following type of input: Image<RGB,224,224>

解决方案

You don't need to do a bunch of image mangling yourself to use a Core ML model with an image — the new Vision framework can do that for you.

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)
    }

}

The WWDC17 session on Vision should have a bit more info — it's tomorrow afternoon.

这篇关于将图像转换为 CVPixelBuffer 用于机器学习 Swift的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！