Project import generated by Copybara.

GitOrigin-RevId: 373e3ac1e5839befd95bf7d73ceff3c5f1171969

docs/solutions/holistic.md

@@ -147,6 +147,18 @@ If set to `true`, the solution filters pose landmarks across different input
 images to reduce jitter, but ignored if [static_image_mode](#static_image_mode)
 is also set to `true`. Default to `true`.
 
+#### enable_segmentation
+
+If set to `true`, in addition to the pose, face and hand landmarks the solution
+also generates the segmentation mask. Default to `false`.
+
+#### smooth_segmentation
+
+If set to `true`, the solution filters segmentation masks across different input
+images to reduce jitter. Ignored if [enable_segmentation](#enable_segmentation)
+is `false` or [static_image_mode](#static_image_mode) is `true`. Default to
+`true`.
+
 #### min_detection_confidence
 
 Minimum confidence value (`[0.0, 1.0]`) from the person-detection model for the
@@ -207,6 +219,15 @@ the camera. The magnitude of `z` uses roughly the same scale as `x`.
 A list of 21 hand landmarks on the right hand, in the same representation as
 [left_hand_landmarks](#left_hand_landmarks).
 
+#### segmentation_mask
+
+The output segmentation mask, predicted only when
+[enable_segmentation](#enable_segmentation) is set to `true`. The mask has the
+same width and height as the input image, and contains values in `[0.0, 1.0]`
+where `1.0` and `0.0` indicate high certainty of a "human" and "background"
+pixel respectively. Please refer to the platform-specific usage examples below
+for usage details.
+
 ### Python Solution API
 
 Please first follow general [instructions](../getting_started/python.md) to
@@ -218,6 +239,8 @@ Supported configuration options:
 *   [static_image_mode](#static_image_mode)
 *   [model_complexity](#model_complexity)
 *   [smooth_landmarks](#smooth_landmarks)
+*   [enable_segmentation](#enable_segmentation)
+*   [smooth_segmentation](#smooth_segmentation)
 *   [min_detection_confidence](#min_detection_confidence)
 *   [min_tracking_confidence](#min_tracking_confidence)
 
@@ -232,7 +255,8 @@ mp_holistic = mp.solutions.holistic
 IMAGE_FILES = []
 with mp_holistic.Holistic(
     static_image_mode=True,
-    model_complexity=2) as holistic:
+    model_complexity=2,
+    enable_segmentation=True) as holistic:
   for idx, file in enumerate(IMAGE_FILES):
     image = cv2.imread(file)
     image_height, image_width, _ = image.shape
@@ -245,8 +269,16 @@ with mp_holistic.Holistic(
           f'{results.pose_landmarks.landmark[mp_holistic.PoseLandmark.NOSE].x * image_width}, '
           f'{results.pose_landmarks.landmark[mp_holistic.PoseLandmark.NOSE].y * image_height})'
       )
-    # Draw pose, left and right hands, and face landmarks on the image.
+
     annotated_image = image.copy()
+    # Draw segmentation on the image.
+    # To improve segmentation around boundaries, consider applying a joint
+    # bilateral filter to "results.segmentation_mask" with "image".
+    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
+    bg_image = np.zeros(image.shape, dtype=np.uint8)
+    bg_image[:] = BG_COLOR
+    annotated_image = np.where(condition, annotated_image, bg_image)
+    # Draw pose, left and right hands, and face landmarks on the image.
     mp_drawing.draw_landmarks(
         annotated_image,
         results.face_landmarks,
@@ -277,12 +309,10 @@ with mp_holistic.Holistic(
       # If loading a video, use 'break' instead of 'continue'.
       continue
 
-    # Flip the image horizontally for a later selfie-view display, and convert
-    # the BGR image to RGB.
-    image = cv2.cvtColor(cv2.flip(image, 1), cv2.COLOR_BGR2RGB)
     # To improve performance, optionally mark the image as not writeable to
     # pass by reference.
     image.flags.writeable = False
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     results = holistic.process(image)
 
     # Draw landmark annotation on the image.
@@ -301,7 +331,8 @@ with mp_holistic.Holistic(
         mp_holistic.POSE_CONNECTIONS,
         landmark_drawing_spec=mp_drawing_styles
         .get_default_pose_landmarks_style())
-    cv2.imshow('MediaPipe Holistic', image)
+    # Flip the image horizontally for a selfie-view display.
+    cv2.imshow('MediaPipe Holistic', cv2.flip(image, 1))
     if cv2.waitKey(5) & 0xFF == 27:
       break
 cap.release()
@@ -317,6 +348,8 @@ Supported configuration options:
 
 *   [modelComplexity](#model_complexity)
 *   [smoothLandmarks](#smooth_landmarks)
+*   [enableSegmentation](#enable_segmentation)
+*   [smoothSegmentation](#smooth_segmentation)
 *   [minDetectionConfidence](#min_detection_confidence)
 *   [minTrackingConfidence](#min_tracking_confidence)
 
@@ -349,8 +382,20 @@ const canvasCtx = canvasElement.getContext('2d');
 function onResults(results) {
   canvasCtx.save();
   canvasCtx.clearRect(0, 0, canvasElement.width, canvasElement.height);
+  canvasCtx.drawImage(results.segmentationMask, 0, 0,
+                      canvasElement.width, canvasElement.height);
+
+  // Only overwrite existing pixels.
+  canvasCtx.globalCompositeOperation = 'source-in';
+  canvasCtx.fillStyle = '#00FF00';
+  canvasCtx.fillRect(0, 0, canvasElement.width, canvasElement.height);
+
+  // Only overwrite missing pixels.
+  canvasCtx.globalCompositeOperation = 'destination-atop';
   canvasCtx.drawImage(
       results.image, 0, 0, canvasElement.width, canvasElement.height);
+
+  canvasCtx.globalCompositeOperation = 'source-over';
   drawConnectors(canvasCtx, results.poseLandmarks, POSE_CONNECTIONS,
                  {color: '#00FF00', lineWidth: 4});
   drawLandmarks(canvasCtx, results.poseLandmarks,
@@ -374,6 +419,8 @@ const holistic = new Holistic({locateFile: (file) => {
 holistic.setOptions({
   modelComplexity: 1,
   smoothLandmarks: true,
+  enableSegmentation: true,
+  smoothSegmentation: true,
   minDetectionConfidence: 0.5,
   minTrackingConfidence: 0.5
 });