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

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

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

Прежде чем вы начнете

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

    pod 'GoogleMLKit/SegmentationSelfie', '3.2.0'

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

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

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

Опции сегментатора

Режим сегментатора

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

STREAM_MODE (default)

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

SINGLE_IMAGE_MODE (default)

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

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

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

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

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

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

Быстрый

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

Цель-C

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

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

Быстрый

let segmenter = Segmenter.segmenter(options: options)

Цель-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

    Цель-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
  }
}

    Цель-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)

    Цель-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.

Цель-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.

Цель-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
}

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