Сегментация селфи с помощью 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. После установки или обновления модулей вашего проекта откройте проект Xcode, используя его . xcworkspace . ML Kit поддерживается в Xcode версии 13.2.1 или выше.

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

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

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

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

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

STREAM_MODE (default)

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

SINGLE_IMAGE_MODE (default)

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

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

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

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

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

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

Быстрый 

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

Objective-C 

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

Наконец, получите экземпляр Segmenter . Передайте указанные вами параметры:

Быстрый 

let segmenter = Segmenter.segmenter(options: options)

Objective-C 

MLKSegmenter *segmenter = [MLKSegmenter segmenterWithOptions:options];

2. Подготовьте входное изображение.

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

Создайте объект VisionImage с помощью UIImage или CMSampleBuffer .

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

  • Создайте объект VisionImage с UIImage . Обязательно укажите правильный .orientation .

    Быстрый 

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

    Objective-C 

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

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

    • Укажите ориентацию данных изображения, содержащихся в CMSampleBuffer .

      Чтобы получить ориентацию изображения:

      Быстрый 

      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 и ориентацию:

      Быстрый 

      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:) .

      Чтобы выполнить сегментацию селфи-изображения синхронно:

      Быстрый 

      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.

      Чтобы выполнить сегментацию селфи-изображения асинхронно:

      Быстрый 

      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. Получите маску сегментации

      Результат сегментации можно получить следующим образом:

      Быстрый 

      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.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:) AVCaptureVideoDataOutput . Оставьте alwaysDiscardsLateVideoFrames AVCaptureVideoDataOutput как true, чтобы ограничить вызовы сегментатора. Если новый видеокадр станет доступен во время работы сегментатора, он будет отброшен.
      • Если вы используете вывод сегментатора для наложения графики на входное изображение, сначала получите результат из ML Kit, затем визуализируйте изображение и наложение за один шаг. При этом визуализируется поверхность отображения только один раз для каждого обработанного входного кадра. См. классы previewOverlayView и CameraViewController в примере быстрого запуска ML Kit для примера.