将图像转换为 CVPixelBuffer 用于机器学习 Swift

Posted 2023-03-27

【中文标题】将图像转换为 CVPixelBuffer 用于机器学习 Swift

【英文标题】：Convert Image to CVPixelBuffer for Machine Learning Swift

【发布时间】：2017-11-08 02:29:47

【问题描述】：

我正在尝试让 Apple 在 2017 年 WWDC 上演示的示例 Core ML 模型正常运行。我正在使用 GoogLeNet 尝试对图像进行分类（请参阅Apple Machine Learning Page）。该模型将 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
    

我从以前的 *** 帖子（最后一个答案 here）中得到了这段代码。我认识到代码可能不正确，但我自己不知道如何做到这一点。我相信这是包含错误的部分。该模型需要以下类型的输入：Image<RGB,224,224>

【问题讨论】：

我创建了一个带有完整代码的示例项目，可以在这里找到：hackernoon.com/…

【参考方案1】：

您无需自己处理大量图像即可使用带有图像的 Core ML 模型 — 新的 Vision framework 可以为您做到这一点。

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)
    

WWDC17 session on Vision 应该有更多信息 - 明天下午。

【讨论】：

像魅力一样工作（有一些修改），谢谢。我没有意识到 Vision 对从图像输入输出信息的模型有特定类型的请求。我想我应该更加关注文档...

对于原问题，VNImageRequestHandler(cgImage: CGImage)更合适。

@chengsam 不是——最初的问题是从磁盘上的资源开始的。将其作为UIImage 读入，转换为CGImage，并将其传递给Vision 会在很长一段时间内丢失元数据，但传递资源URL 会使该元数据对Vision 可用。

如果 MLModel 需要灰度图像，VNImageRequestHandler 是否将其转换为灰度图像？

【参考方案2】：

您可以使用纯 CoreML，但应将图像大小调整为 (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
    

您可以在mimodel 文件中找到输入的预期图像大小：

一个使用纯 CoreML 和 Vision 变体的演示项目，您可以在此处找到：https://github.com/handsomecode/ios11-Demos/tree/coreml_vision/CoreML/CoreMLDemo

【讨论】：

我以为我在 Vision 会话（或者可能是其他 ML 会话之一）中听说您不必调整图像大小......但也许我错了。

@pinkeerach ：如果您使用 Vision API（VNCoreMLRequest，如我的回答），您不必调整图像大小，因为 Vision 会为您处理图像处理部分。如果您直接使用 Core ML（不使用 Vision），则必须调整图像大小并重新格式化（根据您使用的特定模型），然后自己将其转换为 CVPixelBuffer。

@mauryat 您的示例项目什么都不做。真的没有代码。

@zumzum 你可以在这里查看我的示例github.com/handsomecode/iOS11-Demos/tree/coreml_vision，我已经实现了这两种方法

@zumzum 抱歉，我想我没有提交就推动了。在修复之前，我会从 cmets 中删除我的链接。

【参考方案3】：

如果输入的是UIImage，而不是URL，并且你想使用VNImageRequestHandler，你可以使用CIImage。

func updateClassifications(for image: UIImage) 

    let orientation = CGImagePropertyOrientation(image.imageOrientation)

    guard let ciImage = CIImage(image: image) else  return 

    let handler = VNImageRequestHandler(ciImage: ciImage, orientation: orientation)

来自Classifying Images with Vision and Core ML