Project import generated by Copybara.

GitOrigin-RevId: 08c2016a4df5aef571b464a4d4491f38c6b2af10
2021-06-03 17:04:35 -04:00
parent ae05ad04b3
commit 8b57bf879b
118 changed files with 3999 additions and 391 deletions
@@ -0,0 +1,73 @@
+# Copyright 2021 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+load(
+    "//mediapipe/framework/tool:mediapipe_graph.bzl",
+    "mediapipe_simple_subgraph",
+)
+
+licenses(["notice"])
+
+package(default_visibility = ["//visibility:public"])
+
+mediapipe_simple_subgraph(
+    name = "selfie_segmentation_model_loader",
+    graph = "selfie_segmentation_model_loader.pbtxt",
+    register_as = "SelfieSegmentationModelLoader",
+    deps = [
+        "//mediapipe/calculators/core:constant_side_packet_calculator",
+        "//mediapipe/calculators/tflite:tflite_model_calculator",
+        "//mediapipe/calculators/util:local_file_contents_calculator",
+        "//mediapipe/framework/tool:switch_container",
+    ],
+)
+
+mediapipe_simple_subgraph(
+    name = "selfie_segmentation_cpu",
+    graph = "selfie_segmentation_cpu.pbtxt",
+    register_as = "SelfieSegmentationCpu",
+    deps = [
+        ":selfie_segmentation_model_loader",
+        "//mediapipe/calculators/image:image_properties_calculator",
+        "//mediapipe/calculators/tensor:image_to_tensor_calculator",
+        "//mediapipe/calculators/tensor:inference_calculator",
+        "//mediapipe/calculators/tensor:tensors_to_segmentation_calculator",
+        "//mediapipe/calculators/tflite:tflite_custom_op_resolver_calculator",
+        "//mediapipe/calculators/util:from_image_calculator",
+        "//mediapipe/framework/tool:switch_container",
+    ],
+)
+
+mediapipe_simple_subgraph(
+    name = "selfie_segmentation_gpu",
+    graph = "selfie_segmentation_gpu.pbtxt",
+    register_as = "SelfieSegmentationGpu",
+    deps = [
+        ":selfie_segmentation_model_loader",
+        "//mediapipe/calculators/image:image_properties_calculator",
+        "//mediapipe/calculators/tensor:image_to_tensor_calculator",
+        "//mediapipe/calculators/tensor:inference_calculator",
+        "//mediapipe/calculators/tensor:tensors_to_segmentation_calculator",
+        "//mediapipe/calculators/tflite:tflite_custom_op_resolver_calculator",
+        "//mediapipe/calculators/util:from_image_calculator",
+        "//mediapipe/framework/tool:switch_container",
+    ],
+)
+
+exports_files(
+    srcs = [
+        "selfie_segmentation.tflite",
+        "selfie_segmentation_landscape.tflite",
+    ],
+)
@@ -0,0 +1,6 @@
+# selfie_segmentation
+
+Subgraphs|Details
+:--- | :---
+[`SelfieSegmentationCpu`](https://github.com/google/mediapipe/tree/master/mediapipe/modules/selfie_segmentation/selfie_segmentation_cpu.pbtxt)| Segments the person from background in a selfie image. (CPU input, and inference is executed on CPU.)
+[`SelfieSegmentationGpu`](https://github.com/google/mediapipe/tree/master/mediapipe/modules/selfie_segmentation/selfie_segmentation_gpu.pbtxt)| Segments the person from background in a selfie image. (GPU input, and inference is executed on GPU.)
@@ -0,0 +1,131 @@
+# MediaPipe graph to perform selfie segmentation. (CPU input, and all processing
+# and inference are also performed on CPU)
+#
+# It is required that "selfie_segmentation.tflite" or
+# "selfie_segmentation_landscape.tflite" is available at
+# "mediapipe/modules/selfie_segmentation/selfie_segmentation.tflite"
+# or
+# "mediapipe/modules/selfie_segmentation/selfie_segmentation_landscape.tflite"
+# path respectively during execution, depending on the specification in the
+# MODEL_SELECTION input side packet.
+#
+# EXAMPLE:
+#   node {
+#     calculator: "SelfieSegmentationCpu"
+#     input_side_packet: "MODEL_SELECTION:model_selection"
+#     input_stream: "IMAGE:image"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
+#   }
+
+type: "SelfieSegmentationCpu"
+
+# CPU image. (ImageFrame)
+input_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. (ImageFrame in ImageFormat::VEC32F1)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Resizes the input image into a tensor with a dimension desired by the model.
+node {
+  calculator: "SwitchContainer"
+  input_side_packet: "SELECT:model_selection"
+  input_stream: "IMAGE:image"
+  output_stream: "TENSORS:input_tensors"
+  options: {
+    [mediapipe.SwitchContainerOptions.ext] {
+      select: 0
+      contained_node: {
+        calculator: "ImageToTensorCalculator"
+        options: {
+          [mediapipe.ImageToTensorCalculatorOptions.ext] {
+            output_tensor_width: 256
+            output_tensor_height: 256
+            keep_aspect_ratio: false
+            output_tensor_float_range {
+              min: 0.0
+              max: 1.0
+            }
+            border_mode: BORDER_ZERO
+          }
+        }
+      }
+      contained_node: {
+        calculator: "ImageToTensorCalculator"
+        options: {
+          [mediapipe.ImageToTensorCalculatorOptions.ext] {
+            output_tensor_width: 256
+            output_tensor_height: 144
+            keep_aspect_ratio: false
+            output_tensor_float_range {
+              min: 0.0
+              max: 1.0
+            }
+            border_mode: BORDER_ZERO
+          }
+        }
+      }
+    }
+  }
+}
+
+# Generates a single side packet containing a TensorFlow Lite op resolver that
+# supports custom ops needed by the model used in this graph.
+node {
+  calculator: "TfLiteCustomOpResolverCalculator"
+  output_side_packet: "op_resolver"
+}
+
+# Loads the selfie segmentation TF Lite model.
+node {
+  calculator: "SelfieSegmentationModelLoader"
+  input_side_packet: "MODEL_SELECTION:model_selection"
+  output_side_packet: "MODEL:model"
+}
+
+# Runs model inference on CPU.
+node {
+  calculator: "InferenceCalculator"
+  input_stream: "TENSORS:input_tensors"
+  output_stream: "TENSORS:output_tensors"
+  input_side_packet: "MODEL:model"
+  input_side_packet: "CUSTOM_OP_RESOLVER:op_resolver"
+  options: {
+    [mediapipe.InferenceCalculatorOptions.ext] {
+      delegate { xnnpack {} }
+    }
+    #
+  }
+}
+
+# Retrieves the size of the input image.
+node {
+  calculator: "ImagePropertiesCalculator"
+  input_stream: "IMAGE_CPU:image"
+  output_stream: "SIZE:input_size"
+}
+
+# Processes the output tensors into a segmentation mask that has the same size
+# as the input image into the graph.
+node {
+  calculator: "TensorsToSegmentationCalculator"
+  input_stream: "TENSORS:output_tensors"
+  input_stream: "OUTPUT_SIZE:input_size"
+  output_stream: "MASK:mask_image"
+  options: {
+    [mediapipe.TensorsToSegmentationCalculatorOptions.ext] {
+      activation: NONE
+    }
+  }
+}
+
+# Converts the incoming Image into the corresponding ImageFrame type.
+node: {
+  calculator: "FromImageCalculator"
+  input_stream: "IMAGE:mask_image"
+  output_stream: "IMAGE_CPU:segmentation_mask"
+}
@@ -0,0 +1,133 @@
+# MediaPipe graph to perform selfie segmentation. (GPU input, and all processing
+# and inference are also performed on GPU)
+#
+# It is required that "selfie_segmentation.tflite" or
+# "selfie_segmentation_landscape.tflite" is available at
+# "mediapipe/modules/selfie_segmentation/selfie_segmentation.tflite"
+# or
+# "mediapipe/modules/selfie_segmentation/selfie_segmentation_landscape.tflite"
+# path respectively during execution, depending on the specification in the
+# MODEL_SELECTION input side packet.
+#
+# EXAMPLE:
+#   node {
+#     calculator: "SelfieSegmentationGpu"
+#     input_side_packet: "MODEL_SELECTION:model_selection"
+#     input_stream: "IMAGE:image"
+#     output_stream: "SEGMENTATION_MASK:segmentation_mask"
+#   }
+
+type: "SelfieSegmentationGpu"
+
+# GPU image. (GpuBuffer)
+input_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. (GpuBuffer in RGBA, with the same mask values in R and A)
+output_stream: "SEGMENTATION_MASK:segmentation_mask"
+
+# Resizes the input image into a tensor with a dimension desired by the model.
+node {
+  calculator: "SwitchContainer"
+  input_side_packet: "SELECT:model_selection"
+  input_stream: "IMAGE_GPU:image"
+  output_stream: "TENSORS:input_tensors"
+  options: {
+    [mediapipe.SwitchContainerOptions.ext] {
+      select: 0
+      contained_node: {
+        calculator: "ImageToTensorCalculator"
+        options: {
+          [mediapipe.ImageToTensorCalculatorOptions.ext] {
+            output_tensor_width: 256
+            output_tensor_height: 256
+            keep_aspect_ratio: false
+            output_tensor_float_range {
+              min: 0.0
+              max: 1.0
+            }
+            border_mode: BORDER_ZERO
+            gpu_origin: TOP_LEFT
+          }
+        }
+      }
+      contained_node: {
+        calculator: "ImageToTensorCalculator"
+        options: {
+          [mediapipe.ImageToTensorCalculatorOptions.ext] {
+            output_tensor_width: 256
+            output_tensor_height: 144
+            keep_aspect_ratio: false
+            output_tensor_float_range {
+              min: 0.0
+              max: 1.0
+            }
+            border_mode: BORDER_ZERO
+            gpu_origin: TOP_LEFT
+          }
+        }
+      }
+    }
+  }
+}
+
+# Generates a single side packet containing a TensorFlow Lite op resolver that
+# supports custom ops needed by the model used in this graph.
+node {
+  calculator: "TfLiteCustomOpResolverCalculator"
+  output_side_packet: "op_resolver"
+  options: {
+    [mediapipe.TfLiteCustomOpResolverCalculatorOptions.ext] {
+      use_gpu: true
+    }
+  }
+}
+
+# Loads the selfie segmentation TF Lite model.
+node {
+  calculator: "SelfieSegmentationModelLoader"
+  input_side_packet: "MODEL_SELECTION:model_selection"
+  output_side_packet: "MODEL:model"
+}
+
+# Runs model inference on GPU.
+node {
+  calculator: "InferenceCalculator"
+  input_stream: "TENSORS:input_tensors"
+  output_stream: "TENSORS:output_tensors"
+  input_side_packet: "MODEL:model"
+  input_side_packet: "CUSTOM_OP_RESOLVER:op_resolver"
+}
+
+# Retrieves the size of the input image.
+node {
+  calculator: "ImagePropertiesCalculator"
+  input_stream: "IMAGE_GPU:image"
+  output_stream: "SIZE:input_size"
+}
+
+# Processes the output tensors into a segmentation mask that has the same size
+# as the input image into the graph.
+node {
+  calculator: "TensorsToSegmentationCalculator"
+  input_stream: "TENSORS:output_tensors"
+  input_stream: "OUTPUT_SIZE:input_size"
+  output_stream: "MASK:mask_image"
+  options: {
+    [mediapipe.TensorsToSegmentationCalculatorOptions.ext] {
+      activation: NONE
+      gpu_origin: TOP_LEFT
+    }
+  }
+}
+
+# Converts the incoming Image into the corresponding GpuBuffer type.
+node: {
+  calculator: "FromImageCalculator"
+  input_stream: "IMAGE:mask_image"
+  output_stream: "IMAGE_GPU:segmentation_mask"
+}
@@ -0,0 +1,63 @@
+# MediaPipe graph to load a selected selfie segmentation TF Lite model.
+
+type: "SelfieSegmentationModelLoader"
+
+# 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"
+
+# TF Lite model represented as a FlatBuffer.
+# (std::unique_ptr<tflite::FlatBufferModel, std::function<void(tflite::FlatBufferModel*)>>)
+output_side_packet: "MODEL:model"
+
+# Determines path to the desired pose landmark model file.
+node {
+  calculator: "SwitchContainer"
+  input_side_packet: "SELECT:model_selection"
+  output_side_packet: "PACKET:model_path"
+  options: {
+    [mediapipe.SwitchContainerOptions.ext] {
+      select: 0
+      contained_node: {
+        calculator: "ConstantSidePacketCalculator"
+        options: {
+          [mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
+            packet {
+              string_value: "mediapipe/modules/selfie_segmentation/selfie_segmentation.tflite"
+            }
+          }
+        }
+      }
+      contained_node: {
+        calculator: "ConstantSidePacketCalculator"
+        options: {
+          [mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
+            packet {
+              string_value: "mediapipe/modules/selfie_segmentation/selfie_segmentation_landscape.tflite"
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+# Loads the file in the specified path into a blob.
+node {
+  calculator: "LocalFileContentsCalculator"
+  input_side_packet: "FILE_PATH:model_path"
+  output_side_packet: "CONTENTS:model_blob"
+  options: {
+    [mediapipe.LocalFileContentsCalculatorOptions.ext]: {
+      text_mode: false
+    }
+  }
+}
+
+# Converts the input blob into a TF Lite model.
+node {
+  calculator: "TfLiteModelCalculator"
+  input_side_packet: "MODEL_BLOB:model_blob"
+  output_side_packet: "MODEL:model"
+}