將圖像轉換爲用於機器學習的CVPixelBuffer Swift

Question

我正在嘗試在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 
    } 

我從以前的StackOverflow後（最後答案here）得到這個代碼。我意識到代碼可能不正確，但我不知道如何自己做這個。我相信這是包含錯誤的部分。該模型需要以下輸入類型：Image<RGB,224,224>

Answer 1

您不需要做一堆圖像就可以使用核心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應該有更多的信息 - 這是明天下午。

Answer 2

您可以使用純CoreML，但你應該調整到（224224）

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