使用自定义模型给图片加标签 (iOS)

您可以使用机器学习套件识别图片中的实体并为其添加标签。此 API 支持各种自定义图片分类模型。如需有关模型兼容性要求、预训练模型查找位置以及如何训练自有模型的指导,请参阅使用机器学习套件的自定义模型

您可以通过以下两种方式集成自定义模型。您可以将模型放入应用的资源文件夹中以捆绑该模型,也可以从 Cloud Storage 动态下载该模型。下表比较了这两个选项。

捆绑模型 托管模型
模型是应用的 APK 的一部分,这会增加 APK 的大小。 模型不是 APK 的一部分。通过上传到 Cloud Storage 进行托管。我们建议使用 Cloud Storage for Firebase
即使 Android 设备处于离线状态,模型也可立即使用 您的应用必须包含按需下载模型的代码
不需要 Firebase 项目 需要 Firebase 项目(如果使用 Cloud Storage for Firebase)。
您必须重新发布应用才能更新模型 无需重新发布应用即可推送模型更新
没有内置的 A/B 测试 使用 Firebase Remote Config 进行 A/B 测试

试试看

准备工作

  1. 在 Podfile 中添加机器学习套件库:

    pod 'GoogleMLKit/ImageLabelingCustom', '8.0.0'
    
  2. 安装或更新项目的 Pod 之后,请使用 Xcode 项目的 .xcworkspace 来打开项目。Xcode 版本 13.2.1 或更高版本支持机器学习套件。

  3. 如果您想使用 Cloud Storage for Firebase 下载模型,请务必将 Firebase 添加到您的 iOS 项目(如果您尚未添加)。捆绑模型时不需要这样做。

1. 加载模型

配置本地模型来源

如需将模型与您的应用捆绑在一起,请执行以下操作:

  1. 将模型文件(通常以 .tflite.lite 结尾)复制到您的 Xcode 项目,并在执行此操作时注意选择 Copy bundle resources。模型文件将包含在 app bundle 中,并提供给机器学习套件使用。

  2. 创建一个 LocalModel 对象,指定模型文件的路径:

    Swift

    let localModel = LocalModel(path: localModelFilePath)

    Objective-C

    MLKLocalModel *localModel =
        [[MLKLocalModel alloc] initWithPath:localModelFilePath];

配置远程托管的模型来源

如需使用远程托管的模型,您必须使用自己的应用逻辑将模型文件下载到设备的本地存储空间,然后将其作为本地模型加载。我们建议使用 Cloud Storage for Firebase 来托管模型。如需了解实现详情,请参阅 Firebase ML 到 Cloud Storage 迁移指南

配置图片标记器

配置模型来源后,根据其中一个模型创建 ImageLabeler 对象。

提供的选项如下:

选项
confidenceThreshold

检测到的标签的最低置信度得分。如果未设置,系统将使用模型元数据指定的任何分类器阈值。如果模型不包含任何元数据,或者元数据未指定分类器阈值,系统将使用默认阈值 0.0。

maxResultCount

要返回的标签数量上限。如果未设置,系统将使用默认值 10。

如果您只有本地捆绑的模型,只需根据您的 LocalModel 对象创建一个标记器即可:

Swift

let options = CustomImageLabelerOptions(localModel: localModel)
options.confidenceThreshold = NSNumber(value: 0.0)
let imageLabeler = ImageLabeler.imageLabeler(options: options)

Objective-C

MLKCustomImageLabelerOptions *options =
    [[MLKCustomImageLabelerOptions alloc] initWithLocalModel:localModel];
options.confidenceThreshold = @(0.0);
MLKImageLabeler *imageLabeler =
    [MLKImageLabeler imageLabelerWithOptions:options];

如果您使用的是远程托管的模型,则必须在运行之前检查该模型是否已下载。

虽然您只需在运行标签器之前确认这一点,但如果您同时拥有远程托管模型和本地捆绑模型,那么在实例化 ImageLabeler 时执行此检查可能是有意义的:如果远程模型已下载,则从远程模型创建标签器,否则从本地模型创建标签器。

Swift

// Path where your download logic saves the model
let documentDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
let localModelURL = documentDirectory.appendingPathComponent("my_remote_model.tflite")

let model: LocalModel
if FileManager.default.fileExists(atPath: localModelURL.path) {
  // Use the downloaded model
  model = LocalModel(path: localModelURL.path)
} else {
  // Fall back to bundled model
  guard let bundledModelPath = Bundle.main.path(forResource: "model", ofType: "tflite") else { return }
  model = LocalModel(path: bundledModelPath)
}

let options = CustomImageLabelerOptions(localModel: model)
let imageLabeler = ImageLabeler.imageLabeler(options: options)

Objective-C

NSString *documentsDirectory = [NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES) firstObject];
NSString *localModelPath = [documentsDirectory stringByAppendingPathComponent:@"my_remote_model.tflite"];

MLKLocalModel *model;
if ([NSFileManager.defaultManager fileExistsAtPath:localModelPath]) {
  // Use the downloaded model
  model = [[MLKLocalModel alloc] initWithPath:localModelPath];
} else {
  // Fall back to bundled model
  NSString *bundledModelPath = [NSBundle.mainBundle pathForResource:@"model" ofType:@"tflite"];
  model = [[MLKLocalModel alloc] initWithPath:bundledModelPath];
}

MLKCustomImageLabelerOptions *options = [[MLKCustomImageLabelerOptions alloc] initWithLocalModel:model];
MLKImageLabeler *imageLabeler = [MLKImageLabeler imageLabelerWithOptions:options];

如果您只有远程托管的模型,则应停用与模型相关的功能(例如灰显或隐藏部分界面),直到您确认模型已下载。

Swift

let documentDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
let localModelURL = documentDirectory.appendingPathComponent("my_remote_model.tflite")
if FileManager.default.fileExists(atPath: localModelURL.path) {
  // Model is already cached, initialize immediately
  self.initializeLabeler(with: localModelURL)
} else {
  // Model is not yet available, show loading UI and start download
  self.showLoadingUI()
  let storage = Storage.storage()
  let modelRef = storage.reference(forURL: "gs://YOUR_BUCKET/path/to/model.tflite")
  modelRef.write(toFile: localModelURL) { url, error in
    self.hideLoadingUI()
    if let error = error {
      // Handle download error
      self.showErrorUI()
    } else if let modelURL = url {
      // Download success, initialize labeler
      self.initializeLabeler(with: modelURL)
    }
  }
}

func initializeLabeler(with modelURL: URL) {
  let localModel = LocalModel(path: modelURL.path)
  let options = CustomImageLabelerOptions(localModel: localModel)
  self.imageLabeler = ImageLabeler.imageLabeler(options: options)
  // Enable ML-related UI features here
  self.enableMLFeatures()
}

Objective-C

NSString *documentsDirectory = [NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES) firstObject];
NSString *localModelPath = [documentsDirectory stringByAppendingPathComponent:@"my_remote_model.tflite"];
NSURL *localModelURL = [NSURL fileURLWithPath:localModelPath];

if ([NSFileManager.defaultManager fileExistsAtPath:localModelPath]) {
  // Model is already cached, initialize immediately
  [self initializeLabelerWithURL:localModelURL];
} else {
  // Model is not yet available, show loading UI and start download
  [self showLoadingUI];

  FIRStorage *storage = [FIRStorage storage];
  FIRStorageReference *modelRef = [storage referenceForURL:@"gs://YOUR_BUCKET/path/to/model.tflite"];

  [modelRef writeToFile:localModelURL
             completion:^(NSURL * _Nullable URL, NSError * _Nullable error) {
               [self hideLoadingUI];
               if (error != nil) {
                 // Handle download error
                 [self showErrorUI];
               } else {
                 // Download success, initialize labeler
                 [self initializeLabelerWithURL:URL];
               }
             }];
}

- (void)initializeLabelerWithURL:(NSURL *)modelURL {
  MLKLocalModel *localModel = [[MLKLocalModel alloc] initWithPath:modelURL.path];
  MLKCustomImageLabelerOptions *options = [[MLKCustomImageLabelerOptions alloc] initWithLocalModel:localModel];
  self.imageLabeler = [MLKImageLabeler imageLabelerWithOptions:options];

  // Enable ML-related UI features here
  [self enableMLFeatures];
}

2. 准备输入图片

使用 UIImageCMSampleBuffer 创建一个 VisionImage 对象。

如果您使用的是 UIImage,请按以下步骤操作:

  • 使用 UIImage 创建一个 VisionImage 对象。请务必指定正确的 .orientation

    Swift

    let image = VisionImage(image: UIImage)
    visionImage.orientation = image.imageOrientation

    Objective-C

    MLKVisionImage *visionImage = [[MLKVisionImage alloc] initWithImage:image];
    visionImage.orientation = image.imageOrientation;

如果您使用的是 CMSampleBuffer,请按以下步骤操作:

  • 指定 CMSampleBuffer 中所含图片数据的方向。

    如需获取图片方向,请运行以下命令:

    Swift

    func imageOrientation(
      deviceOrientation: UIDeviceOrientation,
      cameraPosition: AVCaptureDevice.Position
    ) -> UIImage.Orientation {
      switch deviceOrientation {
      case .portrait:
        return cameraPosition == .front ? .leftMirrored : .right
      case .landscapeLeft:
        return cameraPosition == .front ? .downMirrored : .up
      case .portraitUpsideDown:
        return cameraPosition == .front ? .rightMirrored : .left
      case .landscapeRight:
        return cameraPosition == .front ? .upMirrored : .down
      case .faceDown, .faceUp, .unknown:
        return .up
      }
    }
          

    Objective-C

    - (UIImageOrientation)
      imageOrientationFromDeviceOrientation:(UIDeviceOrientation)deviceOrientation
                             cameraPosition:(AVCaptureDevicePosition)cameraPosition {
      switch (deviceOrientation) {
        case UIDeviceOrientationPortrait:
          return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationLeftMirrored
                                                                : UIImageOrientationRight;
    
        case UIDeviceOrientationLandscapeLeft:
          return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationDownMirrored
                                                                : UIImageOrientationUp;
        case UIDeviceOrientationPortraitUpsideDown:
          return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationRightMirrored
                                                                : UIImageOrientationLeft;
        case UIDeviceOrientationLandscapeRight:
          return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationUpMirrored
                                                                : UIImageOrientationDown;
        case UIDeviceOrientationUnknown:
        case UIDeviceOrientationFaceUp:
        case UIDeviceOrientationFaceDown:
          return UIImageOrientationUp;
      }
    }
          
  • 使用 CMSampleBuffer 对象和方向创建一个 VisionImage 对象:

    Swift

    let image = VisionImage(buffer: sampleBuffer)
    image.orientation = imageOrientation(
      deviceOrientation: UIDevice.current.orientation,
      cameraPosition: cameraPosition)

    Objective-C

     MLKVisionImage *image = [[MLKVisionImage alloc] initWithBuffer:sampleBuffer];
     image.orientation =
       [self imageOrientationFromDeviceOrientation:UIDevice.currentDevice.orientation
                                    cameraPosition:cameraPosition];

3. 运行图片标记器

如需给图片中的对象加标签,请将 image 对象传递给 ImageLabelerprocess() 方法。

异步:

Swift

imageLabeler.process(image) { labels, error in
    guard error == nil, let labels = labels, !labels.isEmpty else {
        // Handle the error.
        return
    }
    // Show results.
}

Objective-C

[imageLabeler
    processImage:image
      completion:^(NSArray *_Nullable labels,
                   NSError *_Nullable error) {
        if (label.count == 0) {
            // Handle the error.
            return;
        }
        // Show results.
     }];

同步:

Swift

var labels: [ImageLabel]
do {
    labels = try imageLabeler.results(in: image)
} catch let error {
    // Handle the error.
    return
}
// Show results.

Objective-C

NSError *error;
NSArray *labels =
    [imageLabeler resultsInImage:image error:&error];
// Show results or handle the error.

4. 获取有关已加标签的实体的信息

如果为图片添加标签的操作成功,则会返回一组 ImageLabel。每个 ImageLabel 代表在图片中加了标签的某个事物。您可以获取每个标签的文本说明(如果在 LiteRT 模型文件的元数据中可用)、置信度分数和索引。例如:

Swift

for label in labels {
  let labelText = label.text
  let confidence = label.confidence
  let index = label.index
}

Objective-C

for (MLKImageLabel *label in labels) {
  NSString *labelText = label.text;
  float confidence = label.confidence;
  NSInteger index = label.index;
}

提高实时性能的相关提示

如果要在实时应用中给图片加标签,请遵循以下准则以实现最佳帧速率: