在 iOS 系统中使用机器学习套件进行自拍照分割

机器学习套件提供了经过优化的 SDK 来实现自拍分割。自拍分割器资源在构建时静态关联到您的应用。这会将您的应用大小增加多达 24MB,并且 API 延迟时间可能从约 7 毫秒到约 12 毫秒不等,具体取决于输入图像大小(在 iPhone X 上测量)。

试试看

准备工作

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

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

1. 创建 Segmenter 实例

如需对自拍照图片执行分割,请先使用 SelfieSegmenterOptions 创建一个 Segmenter 实例,并视需要指定分割设置。

细分器选项

分割器模式

Segmenter 有两种模式。请务必选择与您的用例相符的设备。

STREAM_MODE (default)

此模式专为流式传输视频或摄像头中的帧而设计。在此模式下,分割器将利用之前帧的结果来返回更流畅的分割结果。

SINGLE_IMAGE_MODE (default)

此模式适用于不相关的单张图片。在此模式下,分割器将独立处理每张图像,而不会对帧进行平滑处理。

启用原始尺寸掩码

要求分割器返回与模型输出大小匹配的原始大小掩码。

原始蒙版尺寸(例如 256x256)通常小于输入图片尺寸。

如果不指定此选项,分割器将重新缩放原始蒙版以匹配输入图像大小。如果您想应用自定义的重新缩放逻辑,或者您的用例不需要重新缩放,请考虑使用此选项。

指定细分器选项:

Swift

let options = SelfieSegmenterOptions()
options.segmenterMode = .singleImage
options.shouldEnableRawSizeMask = true

Objective-C

MLKSelfieSegmenterOptions *options = [[MLKSelfieSegmenterOptions alloc] init];
options.segmenterMode = MLKSegmenterModeSingleImage;
options.shouldEnableRawSizeMask = YES;

最后,获取 Segmenter 的一个实例。传递您指定的选项:

Swift

let segmenter = Segmenter.segmenter(options: options)

Objective-C

MLKSegmenter *segmenter = [MLKSegmenter segmenterWithOptions:options];

2. 准备输入图片

如需对自拍照进行分割,请针对每张图片或视频帧执行以下操作。如果您启用了流模式,则必须从 CMSampleBuffer 创建 VisionImage 对象。

使用 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. 处理图片

VisionImage 对象传递给 Segmenter 的图片处理方法之一。您可以使用异步 process(image:) 方法或同步 results(in:) 方法。

如需同步对自拍照图片执行分割,请执行以下操作:

Swift

var mask: [SegmentationMask]
do {
  mask = try segmenter.results(in: image)
} catch let error {
  print("Failed to perform segmentation with error: \(error.localizedDescription).")
  return
}

// Success. Get a segmentation mask here.

Objective-C

NSError *error;
MLKSegmentationMask *mask =
    [segmenter resultsInImage:image error:&error];
if (error != nil) {
  // Error.
  return;
}

// Success. Get a segmentation mask here.

如需异步对自拍照图片执行分割,请执行以下操作:

Swift

segmenter.process(image) { mask, error in
  guard error == nil else {
    // Error.
    return
  }
  // Success. Get a segmentation mask here.

Objective-C

[segmenter processImage:image
             completion:^(MLKSegmentationMask * _Nullable mask,
                          NSError * _Nullable error) {
               if (error != nil) {
                 // Error.
                 return;
               }
               // Success. Get a segmentation mask here.
             }];

4. 获取分割掩码

您可以获取如下分割结果:

Swift

let maskWidth = CVPixelBufferGetWidth(mask.buffer)
let maskHeight = CVPixelBufferGetHeight(mask.buffer)

CVPixelBufferLockBaseAddress(mask.buffer, CVPixelBufferLockFlags.readOnly)
let maskBytesPerRow = CVPixelBufferGetBytesPerRow(mask.buffer)
var maskAddress =
    CVPixelBufferGetBaseAddress(mask.buffer)!.bindMemory(
        to: Float32.self, capacity: maskBytesPerRow * maskHeight)

for _ in 0...(maskHeight - 1) {
  for col in 0...(maskWidth - 1) {
    // Gets the confidence of the pixel in the mask being in the foreground.
    let foregroundConfidence: Float32 = maskAddress[col]
  }
  maskAddress += maskBytesPerRow / MemoryLayout<Float32>.size
}

Objective-C

size_t width = CVPixelBufferGetWidth(mask.buffer);
size_t height = CVPixelBufferGetHeight(mask.buffer);

CVPixelBufferLockBaseAddress(mask.buffer, kCVPixelBufferLock_ReadOnly);
size_t maskBytesPerRow = CVPixelBufferGetBytesPerRow(mask.buffer);
float *maskAddress = (float *)CVPixelBufferGetBaseAddress(mask.buffer);

for (int row = 0; row < height; ++row) {
  for (int col = 0; col < width; ++col) {
    // Gets the confidence of the pixel in the mask being in the foreground.
    float foregroundConfidence = maskAddress[col];
  }
  maskAddress += maskBytesPerRow / sizeof(float);
}

如需查看如何使用细分结果的完整示例,请参阅机器学习套件快速入门示例

效果提升技巧

结果的质量取决于输入图片的质量:

  • 为了使机器学习套件获得准确的分割结果,图片应至少为 256x256 像素。
  • 如果您在实时应用中执行自拍分割,可能还需要考虑输入图像的整体尺寸。较小的图片的处理速度可能更快,因此,为了缩短延迟时间,请以较低的分辨率捕获图片,但请注意上述分辨率要求,并确保拍摄对象在图片中占据尽可能多的空间。
  • 图片聚焦不佳也会影响准确性。如果您未获得可接受的结果,请让用户重新捕获图片。

如果要在实时应用中使用分割,请遵循以下准则以实现最佳帧速率: