使用 Vision 框架裁剪和缩放照片,并使用 Core ML 模型对其进行分类

使用 Vision 和 Core ML 对图像进行分类

使用 Vision 框架裁剪和缩放照片,并使用 Core ML 模型对其进行分类。

概述

此示例中的应用程序使用 MobileNet 识别图像中最突出的对象,MobileNet 是一种开源图像分类器模型,可识别大约 1,000 个不同的类别。

识别帝王蝶、西兰花和田野中雏菊的应用程序屏幕截图。

每次用户从库中选择照片或使用相机拍摄照片时,应用程序都会将其传递给Vision图像分类请求。Vision 调整照片大小并裁剪照片以满足 MobileNet 模型对其图像输入的约束,然后在幕后使用Core ML框架将照片传递给模型。模型生成预测后,Vision 会将其转发回应用程序,应用程序将结果呈现给用户。

该示例使用 MobileNet 作为如何使用第三方 Core ML 模型的示例。您可以在Core ML 模型库上下载开源模型——包括更新版本的 MobileNet

在您集成第三方模型来解决问题(这可能会增加您的应用程序的大小)之前,请考虑使用 SDK 中的 API。例如,Vision框架的VNClassifyImageRequest类提供与 MobileNet 相同的功能,但可能具有更好的性能并且不会增加应用程序的大小(请参阅分类图像以进行分类和搜索)。

配置示例代码项目

该示例针对 iOS 14 或更高版本,但项目中的 MobileNet 模型适用于:

  • iOS 11 或更高版本
  • macOS 10.13 或更高版本

要在应用程序中拍照,请在带有相机的设备上运行示例。否则,您可以从模拟器中的库中选择照片。

  • 注意:通过将照片拖到模拟器窗口,将您自己的照片添加到模拟器的照片库中。

创建图像分类器实例

在启动时,该类ImagePredictor通过调用其 createImageClassifier()类型方法创建一个图像分类器单例。

/// - Tag: name
static func createImageClassifier() -> VNCoreMLModel {
    // Use a default model configuration.
    let defaultConfig = MLModelConfiguration()

    // Create an instance of the image classifier's wrapper class.
    let imageClassifierWrapper = try? MobileNet(configuration: defaultConfig)

    guard let imageClassifier = imageClassifierWrapper else {
        fatalError("App failed to create an image classifier model instance.")
    }

    // Get the underlying model instance.
    let imageClassifierModel = imageClassifier.model

    // Create a Vision instance using the image classifier's model instance.
    guard let imageClassifierVisionModel = try? VNCoreMLModel(for: imageClassifierModel) else {
        fatalError("App failed to create a `VNCoreMLModel` instance.")
    }

    return imageClassifierVisionModel
}

该方法通过以下方式为 Vision 创建一个 Core ML 模型实例:

  1. 创建 Xcode 在编译时自动生成的模型包装类的实例
  2. 检索包装类实例的底层“MLModel”属性
  3. 将模型实例传递给VNCoreMLModel初始化程序

Image Predictor 类通过仅创建它在应用程序中共享的单个实例来最小化运行时间。

  • VNCoreMLModel注意:为项目中的每个 Core ML 模型共享一个实例。

创建图像分类请求

Image Predictor 类VNCoreMLRequest通过将共享图像分类器模型实例和请求处理程序传递给它的初始化程序来创建一个图像分类请求——一个实例。

// Create an image classification request with an image classifier model.

let imageClassificationRequest = VNCoreMLRequest(model: ImagePredictor.imageClassifier,
                                                 completionHandler: visionRequestHandler)

imageClassificationRequest.imageCropAndScaleOption = .centerCrop

imageCropAndScaleOption该方法通过将请求的属性 设置为 来告诉 Vision 如何调整不符合模型图像输入约束的图像centerCrop

创建请求处理程序

Image Predictor 的方法 通过将图像及其方向传递给初始化器来为每个图像makePredictions(for photo, ...)创建一个 。VNImageRequestHandler

let handler = VNImageRequestHandler(cgImage: photoImage, orientation: orientation)

在将图像发送到模型之前,Vision 会根据实例旋转图像orientationCGImagePropertyOrientation

如果您要分类的图像有一个 URL,请使用以下初始化程序之一创建一个 Vision 图像请求处理程序:

开始请求

该方法通过将请求添加到数组并将其传递给处理程序的方法makePredictions(for photo, ...)来启动请求 VNRequestperform(_:)

let requests: [VNRequest] = [imageClassificationRequest]

// Start the image classification request.
try handler.perform(requests)
  • perform(_:)注意:您可以通过将每个请求添加到您传递给方法参数的数组来对同一图像执行多个视觉请求requests

检索请求的结果

当图像分类请求完成时,Vision 通过调用请求的完成处理程序通知 Image Predictor visionRequestHandler(_:error:)该方法通过以下方式检索请求results

  1. 检查error参数
  2. 转换resultsVNClassificationObservation数组

// Cast the request's results as an `VNClassificationObservation` array.
guard let observations = request.results as? [VNClassificationObservation] else {
    // Image classifiers, like MobileNet, only produce classification observations.
    // However, other Core ML model types can produce other observations.
    // For example, a style transfer model produces `VNPixelBufferObservation` instances.
    print("VNRequest produced the wrong result type: \(type(of: request.results)).")
    return
}

// Create a prediction array from the observations.
predictions = observations.map { observation in
    // Convert each observation into an `ImagePredictor.Prediction` instance.
    Prediction(classification: observation.identifier,
               confidencePercentage: observation.confidencePercentageString)
}

Image Predictor 将每个结果转换为Prediction实例,这是一个具有两个字符串属性的简单结构。

该方法predictions通过调用客户端的完成处理程序将数组发送到 Image Predictor 的客户端——主视图控制器。

// Send the predictions back to the client.
predictionHandler(predictions)

格式化并呈现预测

主视图控制器的imagePredictionHandler(_:)方法将各个预测格式化为单个字符串,并使用辅助方法更新应用程序 UI 中的标签。

private func imagePredictionHandler(_ predictions: [ImagePredictor.Prediction]?) {
    guard let predictions = predictions else {
        updatePredictionLabel("No predictions. (Check console log.)")
        return
    }

    let formattedPredictions = formatPredictions(predictions)

    let predictionString = formattedPredictions.joined(separator: "\n")
    updatePredictionLabel(predictionString)
}

辅助updatePredictionLabel(_:)方法通过更新主调度队列上的标签文本来安全地更新 UI。

func updatePredictionLabel(_ message: String) {
    DispatchQueue.main.async {
        self.predictionLabel.text = message
    }
  • 重要提示:通过在主线程之外使用 Core ML 模型进行预测,让您的应用程序的 UI 保持响应。

GitHub

查看 Github