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

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

Попробуйте!

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

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

    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)

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

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

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

Размер исходной маски (например, 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 .