Сегментация селфи с помощью ML Kit на устройствах iOS

ML Kit предоставляет оптимизированный SDK для сегментации селфи. Объекты Selfie Segmenter статически связываются с приложением во время сборки. Это увеличит размер приложения на 24 МБ, а задержка API может составить от 7 до 12 мс в зависимости от размера входного изображения (измерено на iPhone X).

Попробовать

Подготовка

  1. Добавьте в Podfile следующие библиотеки ML Kit:

    pod 'GoogleMLKit/SegmentationSelfie', '8.0.0'

  2. После установки или обновления Pods откройте проект Xcode, используя его файл .xcworkspace. ML Kit поддерживается в Xcode версии 13.2.1 или более поздней.

1. Создайте экземпляр Segmenter

Чтобы выполнить сегментацию селфи, сначала создайте экземпляр Segmenter с помощью SelfieSegmenterOptions и при необходимости укажите настройки сегментации.

Параметры сегментатора

Режим сегментации

Segmenter работает в двух режимах. Выберите тот, который подходит для вашего случая.

STREAM_MODE (default)

Этот режим предназначен для потоковой передачи кадров из видео или камеры. В этом режиме сегментатор будет использовать результаты предыдущих кадров, чтобы возвращать более плавные результаты сегментации.

SINGLE_IMAGE_MODE (default)

Этот режим предназначен для отдельных изображений, не связанных между собой. В этом режиме сегментатор будет обрабатывать каждое изображение независимо, без сглаживания кадров.

Включить маску исходного размера

Запрашивает у сегментатора возврат маски необработанного размера, которая соответствует размеру выходных данных модели.

Размер необработанной маски (например, 256 x 256) обычно меньше размера входного изображения.

Если не указать этот параметр, сегментатор изменит масштаб необработанной маски в соответствии с размером входного изображения. Используйте этот вариант, если вы хотите применить собственную логику масштабирования или оно вам не требуется.

Укажите параметры сегментатора:

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. Подготовьте входное изображение

Чтобы сегментировать селфи, выполните следующие действия для каждого изображения или кадра видео. Если вы включили режим потоковой передачи, то должны создать объекты VisionImage из CMSampleBuffer.

Создайте объект VisionImage, используя UIImage или CMSampleBuffer.

Если вы используете UIImage, выполните следующие действия:

  • Создайте объект VisionImage с UIImage. Убедитесь, что вы указали правильный .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;
  }
}
  • Создайте объект VisionImage, используя объект CMSampleBuffer и ориентацию:

    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);
}

Полный пример использования результатов сегментации см. в кратком руководстве по ML Kit.

Советы по повышению эффективности

Качество результатов зависит от качества исходного изображения:

  • Чтобы ML Kit получил точный результат сегментации, изображение должно быть размером не менее 256x256 пикселей.
  • Если вы выполняете сегментацию селфи в приложении реального времени, вам также может потребоваться учитывать общие размеры входных изображений. Изображения меньшего размера обрабатываются быстрее, поэтому, чтобы сократить задержку, делайте снимки с более низким разрешением, но помните о требованиях к разрешению, указанных выше, и следите за тем, чтобы объект занимал как можно большую часть изображения.
  • Плохая фокусировка изображения также может повлиять на точность. Если результаты неудовлетворительны, попросите пользователя сделать снимок заново.

Если вы хотите использовать сегментацию в приложении реального времени, следуйте этим рекомендациям, чтобы добиться наилучшей частоты кадров:

  • Используйте режим сегментации stream.
  • Попробуйте сделать снимки с более низким разрешением. Однако не забывайте о требованиях к размерам изображений, которые предъявляет этот API.
  • Для обработки кадров видео используйте синхронный API results(in:) сегментатора. Вызовите этот метод из функции captureOutput(_, didOutput:from:) делегата AVCaptureVideoDataOutputSampleBufferDelegate, чтобы синхронно получить результаты для заданного кадра видео. Для свойства alwaysDiscardsLateVideoFrames объекта AVCaptureVideoDataOutput должно быть задано значение true, чтобы ограничить количество вызовов сегментатора. Если во время работы сегментатора станет доступен новый кадр видео, он будет отброшен.
  • Если вы хотите наложить графику на входное изображение, сначала получите результат от ML Kit, а затем выполните рендеринг изображения и наложение за один шаг. Таким образом, вы выполняете рендеринг на поверхность отображения только один раз для каждого обработанного входного кадра. Пример можно найти в классах previewOverlayView и CameraViewController в кратком руководстве по ML Kit.