初识Core ML：建立一个简单的图片内容识别应用

2017-09-03 本文已影响509人 Andy_Ron

参考：Introduction to Core ML: Building a Simple Image Recognition App

Xcode 9 beta & iOS11 beta & Swift4

Core ML是一个能集成机器学习模型到app中的框架。

一、新建初始项目

新建项目CoreMLDemo，以single-view application template模板
建立如下如下UI

二、实现照相和获取图片库功能

遵守两个协议：
class ViewController: UIViewController, UINavigationControllerDelegate, UIImagePickerControllerDelegate {
添加两个接口，并关联。

  @IBOutlet var imageView: UIImageView!
  @IBOutlet var classifier: UILabel!

实现两个Action

  @IBAction func camera(_ sender: Any) {
      if !UIImagePickerController.isSourceTypeAvailable(.camera) {
          return
      }
      
      let cameraPicker = UIImagePickerController()
      cameraPicker.delegate = self
      cameraPicker.sourceType = .camera
      cameraPicker.allowsEditing = false
      
      present(cameraPicker, animated: true, completion: nil)
  }
  
  @IBAction func openLibrary(_ sender: Any) {
      let picker = UIImagePickerController()
      picker.allowsEditing = false
      picker.delegate = self
      picker.sourceType = .photoLibrary
      present(picker, animated: true)
  }

实现协议UIImagePickerControllerDelegate中的方法：

    func  imagePickerControllerDidCancel(_ picker: UIImagePickerController) {
        dismiss(animated: true, completion: nil)
    }

在Info.plist中添加📷使用请求，Privacy – Camera Usage Description and Privacy

三、集成Core ML Data Model

到Apple机器学习官网 Machine Learning下载Core ML models，目前有6个Model，识别图片的使用Inception v3。下载得到一个mlmodel结尾的文件，直接拖到项目中，会自动生成对应model名的swift的类文件，可以直接在代码中使用。

4.jpg

在ViewController.swift中引入Core ML:
import CoreML
初始Inceptionv3:

    var model: Inceptionv3!
    
    override func viewWillAppear(_ animated: Bool) {
        model = Inceptionv3()
    }

实现UIImagePickerControllerDelegate协议的imagePickerController(_:didFinishPickingMediaWithInfo)方法：

    func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [String : Any]) {
        
        picker.dismiss(animated: true)
        classifier.text = "Analyzing Image..."
        guard let image = info["UIImagePickerControllerOriginalImage"] as? UIImage else {
            return
        }
        
        UIGraphicsBeginImageContextWithOptions(CGSize(width: 299, height: 299), true, 2.0)
        image.draw(in: CGRect(x: 0, y: 0, width: 299, height: 299))
        let newImage = UIGraphicsGetImageFromCurrentImageContext()!
        UIGraphicsEndImageContext()
        
        let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
        var pixelBuffer : CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault, Int(newImage.size.width), Int(newImage.size.height), kCVPixelFormatType_32ARGB, attrs, &pixelBuffer)
        guard (status == kCVReturnSuccess) else {
            return
        }
        
        CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
        let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)
        
        let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
        let context = CGContext(data: pixelData, width: Int(newImage.size.width), height: Int(newImage.size.height), bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue) //3
        
        context?.translateBy(x: 0, y: newImage.size.height)
        context?.scaleBy(x: 1.0, y: -1.0)
        
        UIGraphicsPushContext(context!)
        newImage.draw(in: CGRect(x: 0, y: 0, width: newImage.size.width, height: newImage.size.height))
        UIGraphicsPopContext()
        CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
        imageView.image = newImage

使用Core ML
在上面imagePickerController(_:didFinishPickingMediaWithInfo)方法最后添加：

        guard let prediction = try? model.prediction(image: pixelBuffer!) else {
            return
        }
        classifier.text = "I think this is a \(prediction.classLabel)."

现在就可运行识别图片内容了，我已几个简单的做了实验，效果还可以：

能识别🐶的种类，哈巴狗pug