MLKit: Remove background from video capture using MLKSegmentationMask

Question

I am using MLKit for iOS for Selfie Segmentation. In their sample project they are using a coloured mask to identify the background. I need to remove the background using MLKSegmentationMask and from CVImageBufferRef

https://developers.google.com/ml-kit/vision/selfie-segmentation

Below is the code which takes Segmentation mask from the MLKit and the image buffer which is the actual frame. Now the point is I need to set the background pixel alpha to 0. Segmentation mask contains confidence value from 0 to 1.

+ (void)applySegmentationMask:(MLKSegmentationMask *)mask
                toImageBuffer:(CVImageBufferRef)imageBuffer
          withBackgroundColor:(nullable UIColor *)backgroundColor
              foregroundColor:(nullable UIColor *)foregroundColor {
  NSAssert(CVPixelBufferGetPixelFormatType(imageBuffer) == kCVPixelFormatType_32BGRA,
           @"Image buffer must have 32BGRA pixel format type");
  size_t width = CVPixelBufferGetWidth(mask.buffer);
  size_t height = CVPixelBufferGetHeight(mask.buffer);
  NSAssert(CVPixelBufferGetWidth(imageBuffer) == width, @"Height must match");
  NSAssert(CVPixelBufferGetHeight(imageBuffer) == height, @"Width must match");

  if (backgroundColor == nil && foregroundColor == nil) {
    return;
  }

  CVPixelBufferLockBaseAddress(imageBuffer, 0);
  CVPixelBufferLockBaseAddress(mask.buffer, kCVPixelBufferLock_ReadOnly);

  float *maskAddress = (float *)CVPixelBufferGetBaseAddress(mask.buffer);
  size_t maskBytesPerRow = CVPixelBufferGetBytesPerRow(mask.buffer);

  unsigned char *imageAddress = (unsigned char *)CVPixelBufferGetBaseAddress(imageBuffer);
  size_t bytesPerRow = CVPixelBufferGetBytesPerRow(imageBuffer);
  static const int kBGRABytesPerPixel = 4;

  foregroundColor = foregroundColor ?: UIColor.clearColor;
  backgroundColor = backgroundColor ?: UIColor.clearColor;
  CGFloat redFG, greenFG, blueFG, alphaFG;
  CGFloat redBG, greenBG, blueBG, alphaBG;
  [foregroundColor getRed:&redFG green:&greenFG blue:&blueFG alpha:&alphaFG];
  [backgroundColor getRed:&redBG green:&greenBG blue:&blueBG alpha:&alphaBG];

  static const float kMaxColorComponentValue = 255.0f;

  for (int row = 0; row < height; ++row) {
    for (int col = 0; col < width; ++col) {
      int pixelOffset = col * kBGRABytesPerPixel;
      int blueOffset = pixelOffset;
      int greenOffset = pixelOffset + 1;
      int redOffset = pixelOffset + 2;
      int alphaOffset = pixelOffset + 3;

      float maskValue = maskAddress[col];
      float backgroundRegionRatio = 1.0f - maskValue;
      float foregroundRegionRatio = maskValue;

      float originalPixelRed = imageAddress[redOffset] / kMaxColorComponentValue;
      float originalPixelGreen = imageAddress[greenOffset] / kMaxColorComponentValue;
      float originalPixelBlue = imageAddress[blueOffset] / kMaxColorComponentValue;
      float originalPixelAlpha = imageAddress[alphaOffset] / kMaxColorComponentValue;

      float redOverlay = redBG * backgroundRegionRatio + redFG * foregroundRegionRatio;
      float greenOverlay = greenBG * backgroundRegionRatio + greenFG * foregroundRegionRatio;
      float blueOverlay = blueBG * backgroundRegionRatio + blueFG * foregroundRegionRatio;
      float alphaOverlay = alphaBG * backgroundRegionRatio + alphaFG * foregroundRegionRatio;

      // Calculate composite color component values.
      // Derived from https://en.wikipedia.org/wiki/Alpha_compositing#Alpha_blending
      float compositeAlpha = ((1.0f - alphaOverlay) * originalPixelAlpha) + alphaOverlay;
      float compositeRed = 0.0f;
      float compositeGreen = 0.0f;
      float compositeBlue = 0.0f;
      // Only perform rgb blending calculations if the output alpha is > 0. A zero-value alpha
      // means none of the color channels actually matter, and would introduce division by 0.
      if (fabs(compositeAlpha) > FLT_EPSILON) {
        compositeRed = (((1.0f - alphaOverlay) * originalPixelAlpha * originalPixelRed) +
                        (alphaOverlay * redOverlay)) /
                       compositeAlpha;
        compositeGreen = (((1.0f - alphaOverlay) * originalPixelAlpha * originalPixelGreen) +
                          (alphaOverlay * greenOverlay)) /
                         compositeAlpha;
        compositeBlue = (((1.0f - alphaOverlay) * originalPixelAlpha * originalPixelBlue) +
                         (alphaOverlay * blueOverlay)) /
                        compositeAlpha;
      }

      imageAddress[blueOffset] = compositeBlue * kMaxColorComponentValue;
      imageAddress[greenOffset] = compositeGreen * kMaxColorComponentValue;
      imageAddress[redOffset] = compositeRed * kMaxColorComponentValue;
      imageAddress[alphaOffset] = compositeAlpha * kMaxColorComponentValue;
    }
    imageAddress += bytesPerRow / sizeof(unsigned char);
    maskAddress += maskBytesPerRow / sizeof(float);
  }

  CVPixelBufferUnlockBaseAddress(imageBuffer, 0);
  CVPixelBufferUnlockBaseAddress(mask.buffer, kCVPixelBufferLock_ReadOnly);
}

Answer 1

You haven't posted what you've tried so far, and/or whether what you've tried works. So it's not easy for us to help you, as it seems what you want is someone to do your job.

However, I may suggest looking at the following websites, where they seem to be explaining how to do what you want to achieve:

https://github.com/tbchen/BackgroundRemovalWithCoreMLSample https://medium.com/macoclock/remove-the-image-background-in-swift-using-core-ml-8646ed3a1c14

Regards,