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

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

Попробуйте это

Прежде чем начать

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

    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 .

Создайте объект 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<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 .