Сегментация селфи с помощью 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.