Project import generated by Copybara.

GitOrigin-RevId: 1610e588e497817fae2d9a458093ab6a370e2972
2021-08-18 17:45:46 -07:00
parent b899d17f18
commit 710fb3de58
158 changed files with 10104 additions and 1568 deletions
@@ -101,6 +101,7 @@ mediapipe_simple_subgraph(
    register_as = "FaceLandmarkFrontCpuImage",
    deps = [
        ":face_landmark_front_cpu",
+        "//mediapipe/calculators/image:image_transformation_calculator",
        "//mediapipe/calculators/util:from_image_calculator",
    ],
 )
@@ -111,6 +112,7 @@ mediapipe_simple_subgraph(
    register_as = "FaceLandmarkFrontGpuImage",
    deps = [
        ":face_landmark_front_gpu",
+        "//mediapipe/calculators/image:image_transformation_calculator",
        "//mediapipe/calculators/util:from_image_calculator",
    ],
 )
@@ -8,6 +8,8 @@ input_stream: "IMAGE:image"
 # Max number of faces to detect/track. (int)
 input_side_packet: "NUM_FACES:num_faces"

+# The original input image. (Image)
+output_stream: "IMAGE:image"
 # Collection of detected/predicted faces, each represented as a list of 468 face
 # landmarks. (std::vector<NormalizedLandmarkList>)
 # NOTE: there will not be an output packet in the LANDMARKS stream for this
@@ -30,7 +32,20 @@ output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections"
 node {
  calculator: "FromImageCalculator"
  input_stream: "IMAGE:image"
-  output_stream: "IMAGE_CPU:image_frame"
+  output_stream: "IMAGE_CPU:raw_image_frame"
+  output_stream: "SOURCE_ON_GPU:is_gpu_image"
+}
+
+# TODO: Remove the extra flipping once adopting MlImage.
+# If the source images are on gpu, flip the data vertically before sending them
+# into FaceLandmarkFrontCpu. This maybe needed because OpenGL represents images
+# assuming the image origin is at the bottom-left corner, whereas MediaPipe in
+# general assumes the image origin is at the top-left corner.
+node: {
+  calculator: "ImageTransformationCalculator"
+  input_stream: "IMAGE:raw_image_frame"
+  input_stream: "FLIP_VERTICALLY:is_gpu_image"
+  output_stream: "IMAGE:image_frame"
 }

 node {
@@ -8,6 +8,8 @@ input_stream: "IMAGE:image"
 # Max number of faces to detect/track. (int)
 input_side_packet: "NUM_FACES:num_faces"

+# The original input image. (Image)
+output_stream: "IMAGE:image"
 # Collection of detected/predicted faces, each represented as a list of 468 face
 # landmarks. (std::vector<NormalizedLandmarkList>)
 # NOTE: there will not be an output packet in the LANDMARKS stream for this
@@ -30,6 +32,19 @@ output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections"
 node {
  calculator: "FromImageCalculator"
  input_stream: "IMAGE:image"
+  output_stream: "IMAGE_GPU:raw_gpu_buffer"
+  output_stream: "SOURCE_ON_GPU:is_gpu_image"
+}
+
+# TODO: Remove the extra flipping once adopting MlImage.
+# If the source images are on gpu, flip the data vertically before sending them
+# into FaceLandmarkFrontGpu. This maybe needed because OpenGL represents images
+# assuming the image origin is at the bottom-left corner, whereas MediaPipe in
+# general assumes the image origin is at the top-left corner.
+node: {
+  calculator: "ImageTransformationCalculator"
+  input_stream: "IMAGE_GPU:raw_gpu_buffer"
+  input_stream: "FLIP_VERTICALLY:is_gpu_image"
  output_stream: "IMAGE_GPU:gpu_buffer"
 }

@@ -92,6 +92,7 @@ mediapipe_simple_subgraph(
    register_as = "HandLandmarkTrackingCpuImage",
    deps = [
        ":hand_landmark_tracking_cpu",
+        "//mediapipe/calculators/image:image_transformation_calculator",
        "//mediapipe/calculators/util:from_image_calculator",
    ],
 )
@@ -102,6 +103,7 @@ mediapipe_simple_subgraph(
    register_as = "HandLandmarkTrackingGpuImage",
    deps = [
        ":hand_landmark_tracking_gpu",
+        "//mediapipe/calculators/image:image_transformation_calculator",
        "//mediapipe/calculators/util:from_image_calculator",
    ],
 )
@@ -14,6 +14,8 @@ input_stream: "IMAGE:image"
 # Max number of hands to detect/track. (int)
 input_side_packet: "NUM_HANDS:num_hands"

+# The original input image. (Image)
+output_stream: "IMAGE:image"
 # Collection of detected/predicted hands, each represented as a list of
 # landmarks. (std::vector<NormalizedLandmarkList>)
 # NOTE: there will not be an output packet in the LANDMARKS stream for this
@@ -42,7 +44,20 @@ output_stream: "HAND_ROIS_FROM_PALM_DETECTIONS:hand_rects_from_palm_detections"
 node {
  calculator: "FromImageCalculator"
  input_stream: "IMAGE:image"
-  output_stream: "IMAGE_CPU:image_frame"
+  output_stream: "IMAGE_CPU:raw_image_frame"
+  output_stream: "SOURCE_ON_GPU:is_gpu_image"
+}
+
+# TODO: Remove the extra flipping once adopting MlImage.
+# If the source images are on gpu, flip the data vertically before sending them
+# into HandLandmarkTrackingCpu. This maybe needed because OpenGL represents
+# images assuming the image origin is at the bottom-left corner, whereas
+# MediaPipe in general assumes the image origin is at the top-left corner.
+node: {
+  calculator: "ImageTransformationCalculator"
+  input_stream: "IMAGE:raw_image_frame"
+  input_stream: "FLIP_VERTICALLY:is_gpu_image"
+  output_stream: "IMAGE:image_frame"
 }

 node {
@@ -44,6 +44,19 @@ output_stream: "HAND_ROIS_FROM_PALM_DETECTIONS:hand_rects_from_palm_detections"
 node {
  calculator: "FromImageCalculator"
  input_stream: "IMAGE:image"
+  output_stream: "IMAGE_GPU:raw_gpu_buffer"
+  output_stream: "SOURCE_ON_GPU:is_gpu_image"
+}
+
+# TODO: Remove the extra flipping once adopting MlImage.
+# If the source images are on gpu, flip the data vertically before sending them
+# into HandLandmarkTrackingGpu. This maybe needed because OpenGL represents
+# images assuming the image origin is at the bottom-left corner, whereas
+# MediaPipe in general assumes the image origin is at the top-left corner.
+node: {
+  calculator: "ImageTransformationCalculator"
+  input_stream: "IMAGE_GPU:raw_gpu_buffer"
+  input_stream: "FLIP_VERTICALLY:is_gpu_image"
  output_stream: "IMAGE_GPU:gpu_buffer"
 }

@@ -32,7 +32,9 @@
 #     calculator: "HolisticLandmarkCpu"
 #     input_stream: "IMAGE:input_video"
 #     input_side_packet: "MODEL_COMPLEXITY:model_complexity"
-#     input_side_packet: SMOOTH_LANDMARKS:smooth_landmarks
+#     input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+#     input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
 #     output_stream: "POSE_LANDMARKS:pose_landmarks"
 #     output_stream: "FACE_LANDMARKS:face_landmarks"
 #     output_stream: "LEFT_HAND_LANDMARKS:left_hand_landmarks"
@@ -59,6 +61,14 @@ input_side_packet: "MODEL_COMPLEXITY:model_complexity"
 # If unspecified, functions as set to true. (bool)
 input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"

+# Whether to predict the segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
+# Whether to filter segmentation mask across different input images to reduce
+# jitter. If unspecified, functions as set to true. (bool)
+input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
+
 # Pose landmarks. (NormalizedLandmarkList)
 # 33 pose landmarks.
 output_stream: "POSE_LANDMARKS:pose_landmarks"
@@ -81,6 +91,8 @@ node {
  input_stream: "IMAGE:image"
  input_side_packet: "MODEL_COMPLEXITY:model_complexity"
  input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"
+  input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+  input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
  output_stream: "LANDMARKS:pose_landmarks"
  output_stream: "WORLD_LANDMARKS:pose_world_landmarks"
  output_stream: "ROI_FROM_LANDMARKS:pose_landmarks_roi"
@@ -32,7 +32,9 @@
 #     calculator: "HolisticLandmarkGpu"
 #     input_stream: "IMAGE:input_video"
 #     input_side_packet: "MODEL_COMPLEXITY:model_complexity"
-#     input_side_packet: SMOOTH_LANDMARKS:smooth_landmarks
+#     input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+#     input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
 #     output_stream: "POSE_LANDMARKS:pose_landmarks"
 #     output_stream: "FACE_LANDMARKS:face_landmarks"
 #     output_stream: "LEFT_HAND_LANDMARKS:left_hand_landmarks"
@@ -59,6 +61,14 @@ input_side_packet: "MODEL_COMPLEXITY:model_complexity"
 # If unspecified, functions as set to true. (bool)
 input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"

+# Whether to predict the segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
+# Whether to filter segmentation mask across different input images to reduce
+# jitter. If unspecified, functions as set to true. (bool)
+input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
+
 # Pose landmarks. (NormalizedLandmarkList)
 # 33 pose landmarks.
 output_stream: "POSE_LANDMARKS:pose_landmarks"
@@ -81,6 +91,8 @@ node {
  input_stream: "IMAGE:image"
  input_side_packet: "MODEL_COMPLEXITY:model_complexity"
  input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"
+  input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+  input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
  output_stream: "LANDMARKS:pose_landmarks"
  output_stream: "WORLD_LANDMARKS:pose_world_landmarks"
  output_stream: "ROI_FROM_LANDMARKS:pose_landmarks_roi"
@@ -42,6 +42,10 @@ message AnnotatedKeyPoint {
  int32 id = 1;
  Point3D point_3d = 2;
  NormalizedPoint2D point_2d = 3;
+  // Indicates whether this keypoint is hidden or not. The hidden attribute is
+  // determined from the object's skeleton. For box model, none of the keypoints
+  // are hidden.
+  bool hidden = 4;
 }

 message ObjectAnnotation {
@@ -38,8 +38,10 @@ mediapipe_simple_subgraph(
    graph = "pose_landmark_by_roi_gpu.pbtxt",
    register_as = "PoseLandmarkByRoiGpu",
    deps = [
-        ":pose_landmark_by_roi_post_processing",
        ":pose_landmark_model_loader",
+        ":pose_landmarks_and_segmentation_inverse_projection",
+        ":tensors_to_pose_landmarks_and_segmentation",
+        "//mediapipe/calculators/image:image_properties_calculator",
        "//mediapipe/calculators/tensor:image_to_tensor_calculator",
        "//mediapipe/calculators/tensor:inference_calculator",
    ],
@@ -50,28 +52,41 @@ mediapipe_simple_subgraph(
    graph = "pose_landmark_by_roi_cpu.pbtxt",
    register_as = "PoseLandmarkByRoiCpu",
    deps = [
-        ":pose_landmark_by_roi_post_processing",
        ":pose_landmark_model_loader",
+        ":pose_landmarks_and_segmentation_inverse_projection",
+        ":tensors_to_pose_landmarks_and_segmentation",
+        "//mediapipe/calculators/image:image_properties_calculator",
        "//mediapipe/calculators/tensor:image_to_tensor_calculator",
        "//mediapipe/calculators/tensor:inference_calculator",
    ],
 )

 mediapipe_simple_subgraph(
-    name = "pose_landmark_by_roi_post_processing",
-    graph = "pose_landmark_by_roi_post_processing.pbtxt",
-    register_as = "PoseLandmarkByRoiPostProcessing",
+    name = "tensors_to_pose_landmarks_and_segmentation",
+    graph = "tensors_to_pose_landmarks_and_segmentation.pbtxt",
+    register_as = "TensorsToPoseLandmarksAndSegmentation",
    deps = [
        "//mediapipe/calculators/core:gate_calculator",
        "//mediapipe/calculators/core:split_landmarks_calculator",
        "//mediapipe/calculators/core:split_vector_calculator",
        "//mediapipe/calculators/tensor:tensors_to_floats_calculator",
        "//mediapipe/calculators/tensor:tensors_to_landmarks_calculator",
-        "//mediapipe/calculators/util:landmark_letterbox_removal_calculator",
-        "//mediapipe/calculators/util:landmark_projection_calculator",
+        "//mediapipe/calculators/tensor:tensors_to_segmentation_calculator",
        "//mediapipe/calculators/util:refine_landmarks_from_heatmap_calculator",
        "//mediapipe/calculators/util:thresholding_calculator",
        "//mediapipe/calculators/util:visibility_copy_calculator",
+    ],
+)
+
+mediapipe_simple_subgraph(
+    name = "pose_landmarks_and_segmentation_inverse_projection",
+    graph = "pose_landmarks_and_segmentation_inverse_projection.pbtxt",
+    register_as = "PoseLandmarksAndSegmentationInverseProjection",
+    deps = [
+        "//mediapipe/calculators/image:warp_affine_calculator",
+        "//mediapipe/calculators/util:inverse_matrix_calculator",
+        "//mediapipe/calculators/util:landmark_letterbox_removal_calculator",
+        "//mediapipe/calculators/util:landmark_projection_calculator",
        "//mediapipe/calculators/util:world_landmark_projection_calculator",
    ],
 )
@@ -89,6 +104,17 @@ mediapipe_simple_subgraph(
    ],
 )

+mediapipe_simple_subgraph(
+    name = "pose_segmentation_filtering",
+    graph = "pose_segmentation_filtering.pbtxt",
+    register_as = "PoseSegmentationFiltering",
+    deps = [
+        "//mediapipe/calculators/core:gate_calculator",
+        "//mediapipe/calculators/core:previous_loopback_calculator",
+        "//mediapipe/calculators/image:segmentation_smoothing_calculator",
+    ],
+)
+
 mediapipe_simple_subgraph(
    name = "pose_landmark_gpu",
    graph = "pose_landmark_gpu.pbtxt",
@@ -98,6 +124,7 @@ mediapipe_simple_subgraph(
        ":pose_landmark_by_roi_gpu",
        ":pose_landmark_filtering",
        ":pose_landmarks_to_roi",
+        ":pose_segmentation_filtering",
        "//mediapipe/calculators/core:constant_side_packet_calculator",
        "//mediapipe/calculators/core:gate_calculator",
        "//mediapipe/calculators/core:merge_calculator",
@@ -105,6 +132,7 @@ mediapipe_simple_subgraph(
        "//mediapipe/calculators/core:previous_loopback_calculator",
        "//mediapipe/calculators/core:split_vector_calculator",
        "//mediapipe/calculators/image:image_properties_calculator",
+        "//mediapipe/calculators/util:from_image_calculator",
        "//mediapipe/modules/pose_detection:pose_detection_gpu",
    ],
 )
@@ -118,6 +146,7 @@ mediapipe_simple_subgraph(
        ":pose_landmark_by_roi_cpu",
        ":pose_landmark_filtering",
        ":pose_landmarks_to_roi",
+        ":pose_segmentation_filtering",
        "//mediapipe/calculators/core:constant_side_packet_calculator",
        "//mediapipe/calculators/core:gate_calculator",
        "//mediapipe/calculators/core:merge_calculator",
@@ -125,6 +154,7 @@ mediapipe_simple_subgraph(
        "//mediapipe/calculators/core:previous_loopback_calculator",
        "//mediapipe/calculators/core:split_vector_calculator",
        "//mediapipe/calculators/image:image_properties_calculator",
+        "//mediapipe/calculators/util:from_image_calculator",
        "//mediapipe/modules/pose_detection:pose_detection_cpu",
    ],
 )
@@ -1,5 +1,5 @@
-# MediaPipe graph to detect/predict pose landmarks. (CPU input, and inference is
-# executed on CPU.)
+# MediaPipe graph to detect/predict pose landmarks and optionally segmentation
+# within an ROI. (CPU input, and inference is executed on CPU.)
 #
 # It is required that "pose_landmark_lite.tflite" or
 # "pose_landmark_full.tflite" or "pose_landmark_heavy.tflite" is available at
@@ -13,9 +13,11 @@
 #   node {
 #     calculator: "PoseLandmarkByRoiCpu"
 #     input_side_packet: "MODEL_COMPLEXITY:model_complexity"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
 #     input_stream: "IMAGE:image"
 #     input_stream: "ROI:roi"
 #     output_stream: "LANDMARKS:landmarks"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
 #   }

 type: "PoseLandmarkByRoiCpu"
@@ -26,6 +28,10 @@ input_stream: "IMAGE:image"
 # (NormalizedRect)
 input_stream: "ROI:roi"

+# Whether to predict the segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
 # Complexity of the pose landmark model: 0, 1 or 2. Landmark accuracy as well as
 # inference latency generally go up with the model complexity. If unspecified,
 # functions as set to 1. (int)
@@ -68,7 +74,7 @@ input_side_packet: "MODEL_COMPLEXITY:model_complexity"
 # 31 - left foot index
 # 32 - right foot index
 #
-# NOTE: if a pose is not present within the given ROI, for this particular
+# NOTE: If a pose is not present within the given ROI, for this particular
 # timestamp there will not be an output packet in the LANDMARKS stream. However,
 # the MediaPipe framework will internally inform the downstream calculators of
 # the absence of this packet so that they don't wait for it unnecessarily.
@@ -83,17 +89,30 @@ output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
 # LANDMARKS. However, LANDMARKS provides coordinates (in pixels) of a 3D object
 # projected onto the 2D image surface, while WORLD_LANDMARKS provides
 # coordinates (in meters) of the 3D object itself.
-output_stream: "WORLD_LANDMARKS:pose_world_landmarks"
+output_stream: "WORLD_LANDMARKS:world_landmarks"

-# Transforms the input image into a 256x256 tensor while keeping the aspect
-# ratio (what is expected by the corresponding model), resulting in potential
-# letterboxing in the transformed image.
+# Segmentation mask on CPU in ImageFormat::VEC32F1. (Image)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Retrieves the image size.
+node {
+  calculator: "ImagePropertiesCalculator"
+  input_stream: "IMAGE_CPU:image"
+  output_stream: "SIZE:image_size"
+}
+
+# Crops and transforms the specified ROI in the input image into an image patch
+# represented as a tensor of dimension expected by the corresponding ML model,
+# while maintaining the aspect ratio of the ROI (which can be different from
+# that of the image patch). Therefore, there can be letterboxing around the ROI
+# in the generated tensor representation.
 node: {
  calculator: "ImageToTensorCalculator"
  input_stream: "IMAGE:image"
  input_stream: "NORM_RECT:roi"
  output_stream: "TENSORS:input_tensors"
  output_stream: "LETTERBOX_PADDING:letterbox_padding"
+  output_stream: "MATRIX:transformation_matrix"
  options: {
    [mediapipe.ImageToTensorCalculatorOptions.ext] {
      output_tensor_width: 256
@@ -127,12 +146,33 @@ node {
  }
 }

+# Decodes the tensors into the corresponding landmark and segmentation mask
+# representation.
 node {
-  calculator: "PoseLandmarkByRoiPostProcessing"
+  calculator: "TensorsToPoseLandmarksAndSegmentation"
+  input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
  input_stream: "TENSORS:output_tensors"
-  input_stream: "ROI:roi"
+  output_stream: "LANDMARKS:roi_landmarks"
+  output_stream: "AUXILIARY_LANDMARKS:roi_auxiliary_landmarks"
+  output_stream: "WORLD_LANDMARKS:roi_world_landmarks"
+  output_stream: "SEGMENTATION_MASK:roi_segmentation_mask"
+}
+
+# Projects the landmarks and segmentation mask in the local coordinates of the
+# (potentially letterboxed) ROI back to the global coordinates of the full input
+# image.
+node {
+  calculator: "PoseLandmarksAndSegmentationInverseProjection"
+  input_stream: "IMAGE_SIZE:image_size"
+  input_stream: "NORM_RECT:roi"
  input_stream: "LETTERBOX_PADDING:letterbox_padding"
+  input_stream: "MATRIX:transformation_matrix"
+  input_stream: "LANDMARKS:roi_landmarks"
+  input_stream: "AUXILIARY_LANDMARKS:roi_auxiliary_landmarks"
+  input_stream: "WORLD_LANDMARKS:roi_world_landmarks"
+  input_stream: "SEGMENTATION_MASK:roi_segmentation_mask"
  output_stream: "LANDMARKS:landmarks"
  output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
-  output_stream: "WORLD_LANDMARKS:pose_world_landmarks"
+  output_stream: "WORLD_LANDMARKS:world_landmarks"
+  output_stream: "SEGMENTATION_MASK:segmentation_mask"
 }
@@ -1,5 +1,5 @@
-# MediaPipe graph to detect/predict pose landmarks. (GPU input, and inference is
-# executed on GPU.)
+# MediaPipe graph to detect/predict pose landmarks and optionally segmentation
+# within an ROI. (GPU input, and inference is executed on GPU.)
 #
 # It is required that "pose_landmark_lite.tflite" or
 # "pose_landmark_full.tflite" or "pose_landmark_heavy.tflite" is available at
@@ -13,9 +13,11 @@
 #   node {
 #     calculator: "PoseLandmarkByRoiGpu"
 #     input_side_packet: "MODEL_COMPLEXITY:model_complexity"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
 #     input_stream: "IMAGE:image"
 #     input_stream: "ROI:roi"
 #     output_stream: "LANDMARKS:landmarks"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
 #   }

 type: "PoseLandmarkByRoiGpu"
@@ -26,6 +28,10 @@ input_stream: "IMAGE:image"
 # (NormalizedRect)
 input_stream: "ROI:roi"

+# Whether to predict the segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
 # Complexity of the pose landmark model: 0, 1 or 2. Landmark accuracy as well as
 # inference latency generally go up with the model complexity. If unspecified,
 # functions as set to 1. (int)
@@ -68,7 +74,7 @@ input_side_packet: "MODEL_COMPLEXITY:model_complexity"
 # 31 - left foot index
 # 32 - right foot index
 #
-# NOTE: if a pose is not present within the given ROI, for this particular
+# NOTE: If a pose is not present within the given ROI, for this particular
 # timestamp there will not be an output packet in the LANDMARKS stream. However,
 # the MediaPipe framework will internally inform the downstream calculators of
 # the absence of this packet so that they don't wait for it unnecessarily.
@@ -83,17 +89,30 @@ output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
 # LANDMARKS. However, LANDMARKS provides coordinates (in pixels) of a 3D object
 # projected onto the 2D image surface, while WORLD_LANDMARKS provides
 # coordinates (in meters) of the 3D object itself.
-output_stream: "WORLD_LANDMARKS:pose_world_landmarks"
+output_stream: "WORLD_LANDMARKS:world_landmarks"

-# Transforms the input image into a 256x256 tensor while keeping the aspect
-# ratio (what is expected by the corresponding model), resulting in potential
-# letterboxing in the transformed image.
+# Segmentation mask on GPU in RGBA with the same mask values in R and A. (Image)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Retrieves the image size.
+node {
+  calculator: "ImagePropertiesCalculator"
+  input_stream: "IMAGE_GPU:image"
+  output_stream: "SIZE:image_size"
+}
+
+# Crops and transforms the specified ROI in the input image into an image patch
+# represented as a tensor of dimension expected by the corresponding ML model,
+# while maintaining the aspect ratio of the ROI (which can be different from
+# that of the image patch). Therefore, there can be letterboxing around the ROI
+# in the generated tensor representation.
 node: {
  calculator: "ImageToTensorCalculator"
  input_stream: "IMAGE_GPU:image"
  input_stream: "NORM_RECT:roi"
  output_stream: "TENSORS:input_tensors"
  output_stream: "LETTERBOX_PADDING:letterbox_padding"
+  output_stream: "MATRIX:transformation_matrix"
  options: {
    [mediapipe.ImageToTensorCalculatorOptions.ext] {
      output_tensor_width: 256
@@ -123,12 +142,33 @@ node {
  output_stream: "TENSORS:output_tensors"
 }

+# Decodes the tensors into the corresponding landmark and segmentation mask
+# representation.
 node {
-  calculator: "PoseLandmarkByRoiPostProcessing"
+  calculator: "TensorsToPoseLandmarksAndSegmentation"
+  input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
  input_stream: "TENSORS:output_tensors"
-  input_stream: "ROI:roi"
+  output_stream: "LANDMARKS:roi_landmarks"
+  output_stream: "AUXILIARY_LANDMARKS:roi_auxiliary_landmarks"
+  output_stream: "WORLD_LANDMARKS:roi_world_landmarks"
+  output_stream: "SEGMENTATION_MASK:roi_segmentation_mask"
+}
+
+# Projects the landmarks and segmentation mask in the local coordinates of the
+# (potentially letterboxed) ROI back to the global coordinates of the full input
+# image.
+node {
+  calculator: "PoseLandmarksAndSegmentationInverseProjection"
+  input_stream: "IMAGE_SIZE:image_size"
+  input_stream: "NORM_RECT:roi"
  input_stream: "LETTERBOX_PADDING:letterbox_padding"
+  input_stream: "MATRIX:transformation_matrix"
+  input_stream: "LANDMARKS:roi_landmarks"
+  input_stream: "AUXILIARY_LANDMARKS:roi_auxiliary_landmarks"
+  input_stream: "WORLD_LANDMARKS:roi_world_landmarks"
+  input_stream: "SEGMENTATION_MASK:roi_segmentation_mask"
  output_stream: "LANDMARKS:landmarks"
  output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
-  output_stream: "WORLD_LANDMARKS:pose_world_landmarks"
+  output_stream: "WORLD_LANDMARKS:world_landmarks"
+  output_stream: "SEGMENTATION_MASK:segmentation_mask"
 }
@@ -19,8 +19,11 @@
 #     calculator: "PoseLandmarkCpu"
 #     input_side_packet: "MODEL_COMPLEXITY:model_complexity"
 #     input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+#     input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
 #     input_stream: "IMAGE:image"
 #     output_stream: "LANDMARKS:pose_landmarks"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
 #   }

 type: "PoseLandmarkCpu"
@@ -32,6 +35,14 @@ input_stream: "IMAGE:image"
 # If unspecified, functions as set to true. (bool)
 input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"

+# Whether to predict the segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
+# Whether to filter segmentation mask across different input images to reduce
+# jitter. If unspecified, functions as set to true. (bool)
+input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
+
 # Complexity of the pose landmark model: 0, 1 or 2. Landmark accuracy as well as
 # inference latency generally go up with the model complexity. If unspecified,
 # functions as set to 1. (int)
@@ -88,6 +99,9 @@ output_stream: "LANDMARKS:pose_landmarks"
 # coordinates (in meters) of the 3D object itself.
 output_stream: "WORLD_LANDMARKS:pose_world_landmarks"

+# Segmentation mask. (ImageFrame in ImageFormat::VEC32F1)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
 # Extra outputs (for debugging, for instance).
 # Detected poses. (Detection)
 output_stream: "DETECTION:pose_detection"
@@ -189,11 +203,13 @@ node {
 node {
  calculator: "PoseLandmarkByRoiCpu"
  input_side_packet: "MODEL_COMPLEXITY:model_complexity"
+  input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
  input_stream: "IMAGE:image"
  input_stream: "ROI:pose_rect"
  output_stream: "LANDMARKS:unfiltered_pose_landmarks"
  output_stream: "AUXILIARY_LANDMARKS:unfiltered_auxiliary_landmarks"
  output_stream: "WORLD_LANDMARKS:unfiltered_world_landmarks"
+  output_stream: "SEGMENTATION_MASK:unfiltered_segmentation_mask"
 }

 # Smoothes landmarks to reduce jitter.
@@ -233,3 +249,19 @@ node {
  }
  output_stream: "PREV_LOOP:prev_pose_rect_from_landmarks"
 }
+
+# Smoothes segmentation to reduce jitter.
+node {
+  calculator: "PoseSegmentationFiltering"
+  input_side_packet: "ENABLE:smooth_segmentation"
+  input_stream: "SEGMENTATION_MASK:unfiltered_segmentation_mask"
+  output_stream: "FILTERED_SEGMENTATION_MASK:filtered_segmentation_mask"
+}
+
+# Converts the incoming segmentation mask represented as an Image into the
+# corresponding ImageFrame type.
+node: {
+  calculator: "FromImageCalculator"
+  input_stream: "IMAGE:filtered_segmentation_mask"
+  output_stream: "IMAGE_CPU:segmentation_mask"
+}
@@ -19,8 +19,11 @@
 #     calculator: "PoseLandmarkGpu"
 #     input_side_packet: "MODEL_COMPLEXITY:model_complexity"
 #     input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+#     input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
 #     input_stream: "IMAGE:image"
 #     output_stream: "LANDMARKS:pose_landmarks"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
 #   }

 type: "PoseLandmarkGpu"
@@ -32,6 +35,14 @@ input_stream: "IMAGE:image"
 # If unspecified, functions as set to true. (bool)
 input_side_packet: "SMOOTH_LANDMARKS:smooth_landmarks"

+# Whether to predict the segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
+# Whether to filter segmentation mask across different input images to reduce
+# jitter. If unspecified, functions as set to true. (bool)
+input_side_packet: "SMOOTH_SEGMENTATION:smooth_segmentation"
+
 # Complexity of the pose landmark model: 0, 1 or 2. Landmark accuracy as well as
 # inference latency generally go up with the model complexity. If unspecified,
 # functions as set to 1. (int)
@@ -88,6 +99,9 @@ output_stream: "LANDMARKS:pose_landmarks"
 # coordinates (in meters) of the 3D object itself.
 output_stream: "WORLD_LANDMARKS:pose_world_landmarks"

+# Segmentation mask. (GpuBuffer in RGBA, with the same mask values in R and A)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
 # Extra outputs (for debugging, for instance).
 # Detected poses. (Detection)
 output_stream: "DETECTION:pose_detection"
@@ -189,11 +203,13 @@ node {
 node {
  calculator: "PoseLandmarkByRoiGpu"
  input_side_packet: "MODEL_COMPLEXITY:model_complexity"
+  input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
  input_stream: "IMAGE:image"
  input_stream: "ROI:pose_rect"
  output_stream: "LANDMARKS:unfiltered_pose_landmarks"
  output_stream: "AUXILIARY_LANDMARKS:unfiltered_auxiliary_landmarks"
  output_stream: "WORLD_LANDMARKS:unfiltered_world_landmarks"
+  output_stream: "SEGMENTATION_MASK:unfiltered_segmentation_mask"
 }

 # Smoothes landmarks to reduce jitter.
@@ -233,3 +249,19 @@ node {
  }
  output_stream: "PREV_LOOP:prev_pose_rect_from_landmarks"
 }
+
+# Smoothes segmentation to reduce jitter.
+node {
+  calculator: "PoseSegmentationFiltering"
+  input_side_packet: "ENABLE:smooth_segmentation"
+  input_stream: "SEGMENTATION_MASK:unfiltered_segmentation_mask"
+  output_stream: "FILTERED_SEGMENTATION_MASK:filtered_segmentation_mask"
+}
+
+# Converts the incoming segmentation mask represented as an Image into the
+# corresponding GpuBuffer type.
+node: {
+  calculator: "FromImageCalculator"
+  input_stream: "IMAGE:filtered_segmentation_mask"
+  output_stream: "IMAGE_GPU:segmentation_mask"
+}
@@ -0,0 +1,125 @@
+# MediaPipe graph projecting the landmarks and segmentation mask defined in a
+# local coordinate system within a (potentially letterboxed) ROI back to the
+# global coordinate system of the full image that contains the ROI.
+#
+# EXAMPLE:
+# node {
+#   calculator: "PoseLandmarksAndSegmentationInverseProjection"
+#   input_stream: "IMAGE_SIZE:image_size"
+#   input_stream: "NORM_RECT:roi"
+#   input_stream: "LETTERBOX_PADDING:letterbox_padding"
+#   input_stream: "MATRIX:transformation_matrix"
+#   input_stream: "LANDMARKS:roi_landmarks"
+#   input_stream: "AUXILIARY_LANDMARKS:roi_auxiliary_landmarks"
+#   input_stream: "WORLD_LANDMARKS:roi_world_landmarks"
+#   input_stream: "SEGMENTATION_MASK:roi_segmentation_mask"
+#   output_stream: "LANDMARKS:landmarks"
+#   output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
+#   output_stream: "WORLD_LANDMARKS:world_landmarks"
+#   output_stream: "SEGMENTATION_MASK:segmentation_mask"
+# }
+
+type: "PoseLandmarksAndSegmentationInverseProjection"
+
+# Size of the full image (width & height). (std::pair<int, int>)
+input_stream: "IMAGE_SIZE:image_size"
+
+# ROI within the full image. (NormalizedRect)
+input_stream: "NORM_RECT:roi"
+
+# An array representing the letterbox padding around the ROI from the 4 sides:
+# [left, top, right, bottom]. The padding is normalized to [0.f, 1.f] by the
+# dimensions of the letterboxed/padded ROI. (std::array<float, 4>)
+input_stream: "LETTERBOX_PADDING:letterbox_padding"
+
+# An array representing a 4x4 row-major-order matrix that maps a point within
+# the ROI from the global coordinates of the full image to the local coordinates
+# within the letterboxed ROI. (std::array<float, 16>)
+input_stream: "MATRIX:transformation_matrix"
+
+# Input landmarks and segmentation mask in local coordinates within the
+# letterboxed ROI, and the corresponding outputs in global coordinates of the
+# full image.
+# LANDMARKS & AUXILIARY_LANDMARKS (NormalizedLandmarkList)
+# WORLD_LANDMARKS (LandmarkList)
+# SEGMENTATION_MASK (Image)
+input_stream: "LANDMARKS:roi_landmarks"
+input_stream: "AUXILIARY_LANDMARKS:roi_auxiliary_landmarks"
+input_stream: "WORLD_LANDMARKS:roi_world_landmarks"
+input_stream: "SEGMENTATION_MASK:roi_segmentation_mask"
+output_stream: "LANDMARKS:landmarks"
+output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
+output_stream: "WORLD_LANDMARKS:world_landmarks"
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# -----------------------------------------------------------------------------
+# LANDMARKS
+# -----------------------------------------------------------------------------
+
+# Adjusts landmarks (already normalized to [0.f, 1.f]) in the letterboxed ROI
+# to the corresponding coordinates with the letterbox removed.
+node {
+  calculator: "LandmarkLetterboxRemovalCalculator"
+  input_stream: "LANDMARKS:roi_landmarks"
+  input_stream: "LETTERBOX_PADDING:letterbox_padding"
+  output_stream: "LANDMARKS:adjusted_landmarks"
+}
+node {
+  calculator: "LandmarkLetterboxRemovalCalculator"
+  input_stream: "LANDMARKS:roi_auxiliary_landmarks"
+  input_stream: "LETTERBOX_PADDING:letterbox_padding"
+  output_stream: "LANDMARKS:adjusted_auxiliary_landmarks"
+}
+
+# Projects the landmarks from the letterbox-removed ROI back to the full image.
+node {
+  calculator: "LandmarkProjectionCalculator"
+  input_stream: "NORM_LANDMARKS:adjusted_landmarks"
+  input_stream: "NORM_RECT:roi"
+  output_stream: "NORM_LANDMARKS:landmarks"
+}
+node {
+  calculator: "LandmarkProjectionCalculator"
+  input_stream: "NORM_LANDMARKS:adjusted_auxiliary_landmarks"
+  input_stream: "NORM_RECT:roi"
+  output_stream: "NORM_LANDMARKS:auxiliary_landmarks"
+}
+
+# -----------------------------------------------------------------------------
+# WORLD_LANDMARKS
+# -----------------------------------------------------------------------------
+
+# Projects the world landmarks from the letterboxed ROI to the full image.
+node {
+  calculator: "WorldLandmarkProjectionCalculator"
+  input_stream: "LANDMARKS:roi_world_landmarks"
+  input_stream: "NORM_RECT:roi"
+  output_stream: "LANDMARKS:world_landmarks"
+}
+
+# -----------------------------------------------------------------------------
+# SEGMENTATION_MASK
+# -----------------------------------------------------------------------------
+
+# Calculates the inverse transformation matrix.
+node {
+  calculator: "InverseMatrixCalculator"
+  input_stream: "MATRIX:transformation_matrix"
+  output_stream: "MATRIX:inverse_transformation_matrix"
+}
+
+# Projects the segmentation mask from the letterboxed ROI back to the full
+# image.
+node {
+  calculator: "WarpAffineCalculator"
+  input_stream: "IMAGE:roi_segmentation_mask"
+  input_stream: "MATRIX:inverse_transformation_matrix"
+  input_stream: "OUTPUT_SIZE:image_size"
+  output_stream: "IMAGE:segmentation_mask"
+  options: {
+    [mediapipe.WarpAffineCalculatorOptions.ext] {
+      border_mode: BORDER_ZERO
+      gpu_origin: TOP_LEFT
+    }
+  }
+}
@@ -0,0 +1,61 @@
+# MediaPipe graph to filter segmentation masks temporally (across packets with
+# incremental timestamps) to reduce jitter.
+#
+# EXAMPLE:
+#   node {
+#     calculator: "PoseSegmentationFiltering"
+#     input_side_packet: "ENABLE:enable"
+#     input_stream: "SEGMENTATION_MASK:segmentation_mask"
+#     output_stream: "FILTERED_SEGMENTATION_MASK:filtered_segmentation_mask"
+#   }
+
+type: "PoseSegmentationFiltering"
+
+# Whether to enable filtering. If unspecified, functions as enabled. (bool)
+input_side_packet: "ENABLE:enable"
+
+# Segmentation mask. (Image)
+input_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Filtered segmentation mask. (Image)
+output_stream: "FILTERED_SEGMENTATION_MASK:filtered_segmentation_mask"
+
+# Drops the filtered segmentation mask from the previous frame if filtering is
+# not enabled. In that case, the downstream SegmentationSmoothingCalculator
+# becomes a simple passthrough.
+node {
+  calculator: "GateCalculator"
+  input_side_packet: "ALLOW:enable"
+  input_stream: "prev_filtered_segmentation_mask"
+  output_stream: "gated_prev_filtered_segmentation_mask"
+  options: {
+    [mediapipe.GateCalculatorOptions.ext] {
+      allow: true
+    }
+  }
+}
+
+# Smoothes segmentation to reduce jitter.
+node {
+  calculator: "SegmentationSmoothingCalculator"
+  input_stream: "MASK:segmentation_mask"
+  input_stream: "MASK_PREVIOUS:gated_prev_filtered_segmentation_mask"
+  output_stream: "MASK_SMOOTHED:filtered_segmentation_mask"
+  options {
+    [mediapipe.SegmentationSmoothingCalculatorOptions.ext] {
+      combine_with_previous_ratio: 0.7
+    }
+  }
+}
+
+# Caches the filtered segmentation mask, similar to above for the pose rect.
+node {
+  calculator: "PreviousLoopbackCalculator"
+  input_stream: "MAIN:segmentation_mask"
+  input_stream: "LOOP:filtered_segmentation_mask"
+  input_stream_info: {
+    tag_index: "LOOP"
+    back_edge: true
+  }
+  output_stream: "PREV_LOOP:prev_filtered_segmentation_mask"
+}
@@ -1,34 +1,34 @@
 # MediaPipe graph performing tensor post processing to detect/predict pose
-# landmarks.
+# landmarks and segmenation mask.
 #
 # EXAMPLE:
 #   node {
-#     calculator: "PoseLandmarkByRoiPostProcessing"
+#     calculator: "TensorsToPoseLandmarksAndSegmentation"
+#     input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
 #     input_stream: "TENSORS:tensors"
-#     input_stream: "ROI:roi"
-#     input_stream: "LETTERBOX_PADDING:letterbox_padding"
 #     output_stream: "LANDMARKS:landmarks"
 #     output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
+#     output_stream: "WORLD_LANDMARKS:world_landmarks"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
 #   }

-type: "PoseLandmarkByRoiPostProcessing"
+type: "TensorsToPoseLandmarksAndSegmentation"

-# Tensors from model inference of pose_landmark_{full|heavy|lite}.tflite.
+# Whether to predict segmentation mask. If unspecified, functions as set to
+# false. (bool)
+input_side_packet: "ENABLE_SEGMENTATION:enable_segmentation"
+
+# Tensors from mode inference of
+# "mediapipe/modules/pose_landmark/pose_landmark_lite|full|heavy.tflite".
 # (std::vector<Tensor>)
+# tensors[0]: landmarks
+# tensors[1]: pose flag
+# tensors[2]: segmentation
+# tensors[3]: heatmap
+# tensors[4]: world landmarks
 input_stream: "TENSORS:tensors"

-# ROI (region of interest) within a given image, where only the ROI is used as
-# input to model inference but the resulting pose landmarks are to be reported
-# in coordinates w.r.t. the entire image. (NormalizedRect)
-input_stream: "ROI:roi"
-
-# An array representing the letterbox padding from the 4 sides, i.e.,
-# [left, top, right, bottom], padded around the ROI above before it is used as
-# input to model inference. The padding is normalized to [0.f, 1.f] relatvie to
-# the ROI dimensions. (std::array<float, 4>)
-input_stream: "LETTERBOX_PADDING:letterbox_padding"
-
-# Pose landmarks within the given ROI. (NormalizedLandmarkList)
+# Pose landmarks. (NormalizedLandmarkList)
 # We have 33 landmarks (see pose_landmark_topology.svg) and there are other
 # auxiliary key points.
 # 0 - nose
@@ -65,16 +65,16 @@ input_stream: "LETTERBOX_PADDING:letterbox_padding"
 # 31 - left foot index
 # 32 - right foot index
 #
-# NOTE: if a pose is not present within the given ROI, for this particular
-# timestamp there will not be an output packet in the LANDMARKS stream. However,
-# the MediaPipe framework will internally inform the downstream calculators of
-# the absence of this packet so that they don't wait for it unnecessarily.
+# NOTE: If a pose is not present, for this particular timestamp there will not
+# be an output packet in the LANDMARKS stream. However, the MediaPipe framework
+# will internally inform the downstream calculators of the absence of this
+# packet so that they don't wait for it unnecessarily.
 output_stream: "LANDMARKS:landmarks"
-# Auxiliary landmarks for deriving the ROI in the subsequent image.
+# Auxiliary landmarks (e.g., for deriving the ROI in the subsequent image).
 # (NormalizedLandmarkList)
 output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"

-# Pose world landmarks within the given ROI. (LandmarkList)
+# Pose world landmarks. (LandmarkList)
 # World landmarks are real-world 3D coordinates in meters with the origin at the
 # center between hips. WORLD_LANDMARKS shares the same landmark topology as
 # LANDMARKS. However, LANDMARKS provides coordinates (in pixels) of a 3D object
@@ -82,19 +82,24 @@ output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
 # coordinates (in meters) of the 3D object itself.
 output_stream: "WORLD_LANDMARKS:world_landmarks"

-# Splits a vector of TFLite tensors to multiple vectors according to the ranges
-# specified in option.
+# Segmentation mask. (Image)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Splits a vector of tensors to multiple vectors according to the ranges
+# specified in the option.
 node {
  calculator: "SplitTensorVectorCalculator"
  input_stream: "tensors"
-  output_stream: "landmark_tensors"
+  output_stream: "landmark_tensor"
  output_stream: "pose_flag_tensor"
+  output_stream: "segmentation_tensor"
  output_stream: "heatmap_tensor"
-  output_stream: "world_landmark_tensors"
+  output_stream: "world_landmark_tensor"
  options: {
    [mediapipe.SplitVectorCalculatorOptions.ext] {
      ranges: { begin: 0 end: 1 }
      ranges: { begin: 1 end: 2 }
+      ranges: { begin: 2 end: 3 }
      ranges: { begin: 3 end: 4 }
      ranges: { begin: 4 end: 5 }
    }
@@ -122,21 +127,29 @@ node {
  }
 }

-# Drops landmark tensors if pose is not present.
+# Drops input tensors if pose is not present.
 node {
  calculator: "GateCalculator"
-  input_stream: "landmark_tensors"
-  input_stream: "world_landmark_tensors"
+  input_stream: "landmark_tensor"
+  input_stream: "world_landmark_tensor"
+  input_stream: "segmentation_tensor"
+  input_stream: "heatmap_tensor"
  input_stream: "ALLOW:pose_presence"
-  output_stream: "ensured_landmark_tensors"
-  output_stream: "ensured_world_landmark_tensors"
+  output_stream: "ensured_landmark_tensor"
+  output_stream: "ensured_world_landmark_tensor"
+  output_stream: "ensured_segmentation_tensor"
+  output_stream: "ensured_heatmap_tensor"
 }

+# -----------------------------------------------------------------------------
+# LANDMARKS
+# -----------------------------------------------------------------------------
+
 # Decodes the landmark tensors into a vector of landmarks, where the landmark
-# coordinates are normalized by the size of the input image to the model.
+# coordinates are normalized by the spatial dimensions of the tensor.
 node {
  calculator: "TensorsToLandmarksCalculator"
-  input_stream: "TENSORS:ensured_landmark_tensors"
+  input_stream: "TENSORS:ensured_landmark_tensor"
  output_stream: "NORM_LANDMARKS:raw_landmarks"
  options: {
    [mediapipe.TensorsToLandmarksCalculatorOptions.ext] {
@@ -153,8 +166,8 @@ node {
 node {
  calculator: "RefineLandmarksFromHeatmapCalculator"
  input_stream: "NORM_LANDMARKS:raw_landmarks"
-  input_stream: "TENSORS:heatmap_tensor"
-  output_stream: "NORM_LANDMARKS:refined_landmarks"
+  input_stream: "TENSORS:ensured_heatmap_tensor"
+  output_stream: "NORM_LANDMARKS:all_landmarks"
  options: {
    [mediapipe.RefineLandmarksFromHeatmapCalculatorOptions.ext] {
      kernel_size: 7
@@ -162,26 +175,6 @@ node {
  }
 }

-# Adjusts landmarks (already normalized to [0.f, 1.f]) on the letterboxed pose
-# image (after image transformation with the FIT scale mode) to the
-# corresponding locations on the same image with the letterbox removed (pose
-# image before image transformation).
-node {
-  calculator: "LandmarkLetterboxRemovalCalculator"
-  input_stream: "LANDMARKS:refined_landmarks"
-  input_stream: "LETTERBOX_PADDING:letterbox_padding"
-  output_stream: "LANDMARKS:adjusted_landmarks"
-}
-
-# Projects the landmarks from the cropped pose image to the corresponding
-# locations on the full image before cropping.
-node {
-  calculator: "LandmarkProjectionCalculator"
-  input_stream: "NORM_LANDMARKS:adjusted_landmarks"
-  input_stream: "NORM_RECT:roi"
-  output_stream: "NORM_LANDMARKS:all_landmarks"
-}
-
 # Splits the landmarks into two sets: the actual pose landmarks and the
 # auxiliary landmarks.
 node {
@@ -197,11 +190,15 @@ node {
  }
 }

-# Decodes the world landmark tensors into a vector of world landmarks.
+# -----------------------------------------------------------------------------
+# WORLD_LANDMARKS
+# -----------------------------------------------------------------------------
+
+# Decodes the world-landmark tensors into a vector of world landmarks.
 node {
  calculator: "TensorsToLandmarksCalculator"
-  input_stream: "TENSORS:ensured_world_landmark_tensors"
-  output_stream: "LANDMARKS:raw_world_landmarks"
+  input_stream: "TENSORS:ensured_world_landmark_tensor"
+  output_stream: "LANDMARKS:all_world_landmarks"
  options: {
    [mediapipe.TensorsToLandmarksCalculatorOptions.ext] {
      num_landmarks: 39
@@ -209,13 +206,25 @@ node {
  }
 }

+# Keeps only the actual world landmarks.
+node {
+  calculator: "SplitLandmarkListCalculator"
+  input_stream: "all_world_landmarks"
+  output_stream: "world_landmarks_without_visibility"
+  options: {
+    [mediapipe.SplitVectorCalculatorOptions.ext] {
+      ranges: { begin: 0 end: 33 }
+    }
+  }
+}
+
 # Reuses the visibility and presence field in pose landmarks for the world
 # landmarks.
 node {
  calculator: "VisibilityCopyCalculator"
-  input_stream: "NORM_LANDMARKS_FROM:raw_landmarks"
-  input_stream: "LANDMARKS_TO:raw_world_landmarks"
-  output_stream: "LANDMARKS_TO:raw_world_landmarks_with_visibility"
+  input_stream: "NORM_LANDMARKS_FROM:landmarks"
+  input_stream: "LANDMARKS_TO:world_landmarks_without_visibility"
+  output_stream: "LANDMARKS_TO:world_landmarks"
  options: {
    [mediapipe.VisibilityCopyCalculatorOptions.ext] {
      copy_visibility: true
@@ -224,23 +233,33 @@ node {
  }
 }

-# Projects the world landmarks from the cropped pose image to the corresponding
-# locations on the full image before cropping.
-node {
-  calculator: "WorldLandmarkProjectionCalculator"
-  input_stream: "LANDMARKS:raw_world_landmarks_with_visibility"
-  input_stream: "NORM_RECT:roi"
-  output_stream: "LANDMARKS:all_world_landmarks"
-}
+# -----------------------------------------------------------------------------
+# SEGMENTATION_MASK
+# -----------------------------------------------------------------------------

-# Takes only actual world landmarks.
+# Drops segmentation tensors if segmentation is not enabled.
 node {
-  calculator: "SplitLandmarkListCalculator"
-  input_stream: "all_world_landmarks"
-  output_stream: "world_landmarks"
+  calculator: "GateCalculator"
+  input_side_packet: "ALLOW:enable_segmentation"
+  input_stream: "ensured_segmentation_tensor"
+  output_stream: "enabled_segmentation_tensor"
  options: {
-    [mediapipe.SplitVectorCalculatorOptions.ext] {
-      ranges: { begin: 0 end: 33 }
+    [mediapipe.GateCalculatorOptions.ext] {
+      allow: true
+    }
+  }
+}
+
+# Decodes the segmentation tensor into a mask image with pixel values in [0, 1]
+# (1 for person and 0 for background).
+node {
+  calculator: "TensorsToSegmentationCalculator"
+  input_stream: "TENSORS:enabled_segmentation_tensor"
+  output_stream: "MASK:segmentation_mask"
+  options: {
+    [mediapipe.TensorsToSegmentationCalculatorOptions.ext] {
+      activation: SIGMOID
+      gpu_origin: TOP_LEFT
    }
  }
 }
@@ -65,6 +65,30 @@ mediapipe_simple_subgraph(
    ],
 )

+mediapipe_simple_subgraph(
+    name = "selfie_segmentation_cpu_image",
+    graph = "selfie_segmentation_cpu_image.pbtxt",
+    register_as = "SelfieSegmentationCpuImage",
+    deps = [
+        ":selfie_segmentation_cpu",
+        "//mediapipe/calculators/image:image_transformation_calculator",
+        "//mediapipe/calculators/util:from_image_calculator",
+        "//mediapipe/calculators/util:to_image_calculator",
+    ],
+)
+
+mediapipe_simple_subgraph(
+    name = "selfie_segmentation_gpu_image",
+    graph = "selfie_segmentation_gpu_image.pbtxt",
+    register_as = "SelfieSegmentationGpuImage",
+    deps = [
+        ":selfie_segmentation_gpu",
+        "//mediapipe/calculators/image:image_transformation_calculator",
+        "//mediapipe/calculators/util:from_image_calculator",
+        "//mediapipe/calculators/util:to_image_calculator",
+    ],
+)
+
 exports_files(
    srcs = [
        "selfie_segmentation.tflite",
@@ -0,0 +1,50 @@
+# MediaPipe graph to perform selfie segmentation.
+
+type: "SelfieSegmentationCpuImage"
+
+# Input image. (Image)
+input_stream: "IMAGE:image"
+
+# The original input image. (Image)
+output_stream: "IMAGE:image"
+
+# An integer 0 or 1. Use 0 to select a general-purpose model (operating on a
+# 256x256 tensor), and 1 to select a model (operating on a 256x144 tensor) more
+# optimized for landscape images. If unspecified, functions as set to 0. (int)
+input_side_packet: "MODEL_SELECTION:model_selection"
+
+# Segmentation mask. (Image)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Converts Image to ImageFrame for SelfieSegmentationCpu to consume.
+node {
+  calculator: "FromImageCalculator"
+  input_stream: "IMAGE:image"
+  output_stream: "IMAGE_CPU:raw_image_frame"
+  output_stream: "SOURCE_ON_GPU:is_gpu_image"
+}
+
+# TODO: Remove the extra flipping once adopting MlImage.
+# If the source images are on gpu, flip the data vertically before sending them
+# into SelfieSegmentationCpu. This maybe needed because OpenGL represents images
+# assuming the image origin is at the bottom-left corner, whereas MediaPipe in
+# general assumes the image origin is at the top-left corner.
+node: {
+  calculator: "ImageTransformationCalculator"
+  input_stream: "IMAGE:raw_image_frame"
+  input_stream: "FLIP_VERTICALLY:is_gpu_image"
+  output_stream: "IMAGE:image_frame"
+}
+
+node {
+  calculator: "SelfieSegmentationCpu"
+  input_side_packet: "MODEL_SELECTION:model_selection"
+  input_stream: "IMAGE:image_frame"
+  output_stream: "SEGMENTATION_MASK:segmentation_mask_image_frame"
+}
+
+node {
+  calculator: "ToImageCalculator"
+  input_stream: "IMAGE_CPU:segmentation_mask_image_frame"
+  output_stream: "IMAGE:segmentation_mask"
+}
@@ -0,0 +1,50 @@
+# MediaPipe graph to perform selfie segmentation.
+
+type: "SelfieSegmentationGpuImage"
+
+# Input image. (Image)
+input_stream: "IMAGE:image"
+
+# The original input image. (Image)
+output_stream: "IMAGE:image"
+
+# An integer 0 or 1. Use 0 to select a general-purpose model (operating on a
+# 256x256 tensor), and 1 to select a model (operating on a 256x144 tensor) more
+# optimized for landscape images. If unspecified, functions as set to 0. (int)
+input_side_packet: "MODEL_SELECTION:model_selection"
+
+# Segmentation mask. (Image)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Converts Image to ImageFrame for SelfieSegmentationGpu to consume.
+node {
+  calculator: "FromImageCalculator"
+  input_stream: "IMAGE:image"
+  output_stream: "IMAGE_GPU:raw_gpu_buffer"
+  output_stream: "SOURCE_ON_GPU:is_gpu_image"
+}
+
+# TODO: Remove the extra flipping once adopting MlImage.
+# If the source images are on gpu, flip the data vertically before sending them
+# into SelfieSegmentationGpu. This maybe needed because OpenGL represents images
+# assuming the image origin is at the bottom-left corner, whereas MediaPipe in
+# general assumes the image origin is at the top-left corner.
+node: {
+  calculator: "ImageTransformationCalculator"
+  input_stream: "IMAGE_GPU:raw_gpu_buffer"
+  input_stream: "FLIP_VERTICALLY:is_gpu_image"
+  output_stream: "IMAGE_GPU:gpu_buffer"
+}
+
+node {
+  calculator: "SelfieSegmentationGpu"
+  input_side_packet: "MODEL_SELECTION:model_selection"
+  input_stream: "IMAGE:gpu_buffer"
+  output_stream: "SEGMENTATION_MASK:segmentation_mask_gpu_buffer"
+}
+
+node {
+  calculator: "ToImageCalculator"
+  input_stream: "IMAGE_GPU:segmentation_mask_gpu_buffer"
+  output_stream: "IMAGE:segmentation_mask"
+}