From da02052c70b138d6d5ea05d51103e02c535a7542 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Wed, 27 Sep 2023 17:11:09 -0700
Subject: [PATCH] landmarks_to_tensor stream utility function.

PiperOrigin-RevId: 569003241
---
 mediapipe/framework/api2/stream/BUILD         | 32 +++++++
 .../api2/stream/landmarks_to_tensor.cc        | 64 +++++++++++++
 .../api2/stream/landmarks_to_tensor.h         | 37 ++++++++
 .../api2/stream/landmarks_to_tensor_test.cc   | 89 +++++++++++++++++++
 4 files changed, 222 insertions(+)
 create mode 100644 mediapipe/framework/api2/stream/landmarks_to_tensor.cc
 create mode 100644 mediapipe/framework/api2/stream/landmarks_to_tensor.h
 create mode 100644 mediapipe/framework/api2/stream/landmarks_to_tensor_test.cc

diff --git a/mediapipe/framework/api2/stream/BUILD b/mediapipe/framework/api2/stream/BUILD
index 88607844e..323c36713 100644
--- a/mediapipe/framework/api2/stream/BUILD
+++ b/mediapipe/framework/api2/stream/BUILD
@@ -123,6 +123,38 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "landmarks_to_tensor",
+    srcs = ["landmarks_to_tensor.cc"],
+    hdrs = ["landmarks_to_tensor.h"],
+    deps = [
+        "//mediapipe/calculators/tensor:landmarks_to_tensor_calculator",
+        "//mediapipe/calculators/tensor:landmarks_to_tensor_calculator_cc_proto",
+        "//mediapipe/framework/api2:builder",
+        "//mediapipe/framework/api2:port",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/formats:tensor",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "landmarks_to_tensor_test",
+    srcs = ["landmarks_to_tensor_test.cc"],
+    deps = [
+        ":landmarks_to_tensor",
+        "//mediapipe/calculators/tensor:landmarks_to_tensor_calculator_cc_proto",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/api2:builder",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/formats:tensor",
+        "//mediapipe/framework/port:gtest",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/port:parse_text_proto",
+        "//mediapipe/framework/port:status_matchers",
+    ],
+)
+
 cc_library(
     name = "landmarks_projection",
     srcs = ["landmarks_projection.cc"],
diff --git a/mediapipe/framework/api2/stream/landmarks_to_tensor.cc b/mediapipe/framework/api2/stream/landmarks_to_tensor.cc
new file mode 100644
index 000000000..85ce7a04b
--- /dev/null
+++ b/mediapipe/framework/api2/stream/landmarks_to_tensor.cc
@@ -0,0 +1,64 @@
+#include "mediapipe/framework/api2/stream/landmarks_to_tensor.h"
+
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "mediapipe/calculators/tensor/landmarks_to_tensor_calculator.h"
+#include "mediapipe/framework/api2/builder.h"
+#include "mediapipe/framework/api2/port.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
+#include "mediapipe/framework/formats/tensor.h"
+
+namespace mediapipe::api2::builder {
+
+namespace {
+
+using ::mediapipe::api2::LandmarksToTensorCalculator;
+
+template <typename LandmarkListType>
+Stream<std::vector<Tensor>> InternalConvertToTensor(
+    Stream<LandmarkListType> landmarks,
+    std::optional<Stream<std::pair<int, int>>> image_size,
+    absl::Span<const LandmarksToTensorCalculatorOptions::Attribute> attributes,
+    const bool flatten, Graph& graph) {
+  auto& to_tensor = graph.AddNode<LandmarksToTensorCalculator>();
+  auto& to_tensor_options =
+      to_tensor.GetOptions<LandmarksToTensorCalculatorOptions>();
+  for (const auto& attribute : attributes) {
+    to_tensor_options.add_attributes(attribute);
+  }
+  to_tensor_options.set_flatten(flatten);
+  if constexpr (std::is_same_v<LandmarkListType, LandmarkList>) {
+    landmarks.ConnectTo(
+        to_tensor[LandmarksToTensorCalculator::kInLandmarkList]);
+  } else {
+    landmarks.ConnectTo(
+        to_tensor[LandmarksToTensorCalculator::kInNormLandmarkList]);
+  }
+  if (image_size.has_value()) {
+    image_size->ConnectTo(to_tensor[LandmarksToTensorCalculator::kImageSize]);
+  }
+  return to_tensor[LandmarksToTensorCalculator::kOutTensors];
+}
+
+}  // namespace
+
+Stream<std::vector<Tensor>> ConvertLandmarksToTensor(
+    Stream<LandmarkList> landmarks,
+    absl::Span<const LandmarksToTensorCalculatorOptions::Attribute> attributes,
+    const bool flatten, Graph& graph) {
+  return InternalConvertToTensor(landmarks, /*image_size=*/std::nullopt,
+                                 attributes, flatten, graph);
+}
+
+Stream<std::vector<Tensor>> ConvertNormalizedLandmarksToTensor(
+    Stream<NormalizedLandmarkList> normalized_landmarks,
+    Stream<std::pair<int, int>> image_size,
+    absl::Span<const LandmarksToTensorCalculatorOptions::Attribute> attributes,
+    const bool flatten, Graph& graph) {
+  return InternalConvertToTensor(normalized_landmarks, image_size, attributes,
+                                 flatten, graph);
+}
+
+}  // namespace mediapipe::api2::builder
diff --git a/mediapipe/framework/api2/stream/landmarks_to_tensor.h b/mediapipe/framework/api2/stream/landmarks_to_tensor.h
new file mode 100644
index 000000000..3766c53d2
--- /dev/null
+++ b/mediapipe/framework/api2/stream/landmarks_to_tensor.h
@@ -0,0 +1,37 @@
+#ifndef MEDIAPIPE_FRAMEWORK_API2_STREAM_LANDMARKS_TO_TENSOR_H_
+#define MEDIAPIPE_FRAMEWORK_API2_STREAM_LANDMARKS_TO_TENSOR_H_
+
+#include <utility>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "mediapipe/calculators/tensor/landmarks_to_tensor_calculator.pb.h"
+#include "mediapipe/framework/api2/builder.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
+#include "mediapipe/framework/formats/tensor.h"
+
+namespace mediapipe::api2::builder {
+
+// Updates @graph to convert @landmarks to a Tensor. Values and their order are
+// defined by @attributes. If @flatten is true resulting tensor will be 1D,
+// otherwise tensor will be 2D with (n_landmarks, n_attributes) shape.
+Stream<std::vector<Tensor>> ConvertLandmarksToTensor(
+    Stream<mediapipe::LandmarkList> landmarks,
+    absl::Span<const mediapipe::LandmarksToTensorCalculatorOptions::Attribute>
+        attributes,
+    bool flatten, Graph& graph);
+
+// Updates @graph to convert @normalized_landmarks to a Tensor. Values and their
+// order are defined by @attributes. X, Y and Z values are scaled using
+// @image_size. If @flatten is true resulting tensor will be 1D, otherwise
+// tensor will be 2D with (n_landmarks, n_attributes) shape.
+Stream<std::vector<Tensor>> ConvertNormalizedLandmarksToTensor(
+    Stream<mediapipe::NormalizedLandmarkList> normalized_landmarks,
+    Stream<std::pair<int, int>> image_size,
+    absl::Span<const mediapipe::LandmarksToTensorCalculatorOptions::Attribute>
+        attributes,
+    bool flatten, Graph& graph);
+
+}  // namespace mediapipe::api2::builder
+
+#endif  // MEDIAPIPE_FRAMEWORK_API2_STREAM_LANDMARKS_TO_TENSOR_H_
diff --git a/mediapipe/framework/api2/stream/landmarks_to_tensor_test.cc b/mediapipe/framework/api2/stream/landmarks_to_tensor_test.cc
new file mode 100644
index 000000000..a4fba3e5e
--- /dev/null
+++ b/mediapipe/framework/api2/stream/landmarks_to_tensor_test.cc
@@ -0,0 +1,89 @@
+#include "mediapipe/framework/api2/stream/landmarks_to_tensor.h"
+
+#include <utility>
+#include <vector>
+
+#include "mediapipe/calculators/tensor/landmarks_to_tensor_calculator.pb.h"
+#include "mediapipe/framework/api2/builder.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
+#include "mediapipe/framework/formats/tensor.h"
+#include "mediapipe/framework/port/gmock.h"
+#include "mediapipe/framework/port/gtest.h"
+#include "mediapipe/framework/port/parse_text_proto.h"
+#include "mediapipe/framework/port/status_matchers.h"
+
+namespace mediapipe::api2::builder {
+namespace {
+
+TEST(ConvertLandmarksToTensor, ConvertLandmarksToTensor) {
+  Graph graph;
+
+  Stream<LandmarkList> landmarks = graph.In("LANDMARKS").Cast<LandmarkList>();
+  Stream<std::vector<Tensor>> tensors =
+      ConvertLandmarksToTensor(landmarks,
+                               {LandmarksToTensorCalculatorOptions::X,
+                                LandmarksToTensorCalculatorOptions::Y,
+                                LandmarksToTensorCalculatorOptions::Z},
+                               /*flatten=*/true, graph);
+  tensors.SetName("tensors");
+
+  EXPECT_THAT(graph.GetConfig(),
+              EqualsProto(ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+                node {
+                  calculator: "LandmarksToTensorCalculator"
+                  input_stream: "LANDMARKS:__stream_0"
+                  output_stream: "TENSORS:tensors"
+                  options {
+                    [mediapipe.LandmarksToTensorCalculatorOptions.ext] {
+                      attributes: [ X, Y, Z ]
+                      flatten: true
+                    }
+                  }
+                }
+                input_stream: "LANDMARKS:__stream_0"
+              )pb")));
+
+  CalculatorGraph calcualtor_graph;
+  MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
+}
+
+TEST(ConvertLandmarksToTensor, ConvertNormalizedLandmarksToTensor) {
+  Graph graph;
+
+  Stream<NormalizedLandmarkList> landmarks =
+      graph.In("LANDMARKS").Cast<NormalizedLandmarkList>();
+  Stream<std::pair<int, int>> image_size =
+      graph.In("IMAGE_SIZE").Cast<std::pair<int, int>>();
+  Stream<std::vector<Tensor>> tensors = ConvertNormalizedLandmarksToTensor(
+      landmarks, image_size,
+      {LandmarksToTensorCalculatorOptions::X,
+       LandmarksToTensorCalculatorOptions::Y,
+       LandmarksToTensorCalculatorOptions::Z},
+      /*flatten=*/false, graph);
+  tensors.SetName("tensors");
+
+  EXPECT_THAT(graph.GetConfig(),
+              EqualsProto(ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+                node {
+                  calculator: "LandmarksToTensorCalculator"
+                  input_stream: "IMAGE_SIZE:__stream_0"
+                  input_stream: "NORM_LANDMARKS:__stream_1"
+                  output_stream: "TENSORS:tensors"
+                  options {
+                    [mediapipe.LandmarksToTensorCalculatorOptions.ext] {
+                      attributes: [ X, Y, Z ]
+                      flatten: false
+                    }
+                  }
+                }
+                input_stream: "IMAGE_SIZE:__stream_0"
+                input_stream: "LANDMARKS:__stream_1"
+              )pb")));
+
+  CalculatorGraph calcualtor_graph;
+  MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
+}
+
+}  // namespace
+}  // namespace mediapipe::api2::builder