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 +#include +#include + +#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 +Stream> InternalConvertToTensor( + Stream landmarks, + std::optional>> image_size, + absl::Span attributes, + const bool flatten, Graph& graph) { + auto& to_tensor = graph.AddNode(); + auto& to_tensor_options = + to_tensor.GetOptions(); + for (const auto& attribute : attributes) { + to_tensor_options.add_attributes(attribute); + } + to_tensor_options.set_flatten(flatten); + if constexpr (std::is_same_v) { + 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> ConvertLandmarksToTensor( + Stream landmarks, + absl::Span attributes, + const bool flatten, Graph& graph) { + return InternalConvertToTensor(landmarks, /*image_size=*/std::nullopt, + attributes, flatten, graph); +} + +Stream> ConvertNormalizedLandmarksToTensor( + Stream normalized_landmarks, + Stream> image_size, + absl::Span 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 +#include + +#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> ConvertLandmarksToTensor( + Stream landmarks, + absl::Span + 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> ConvertNormalizedLandmarksToTensor( + Stream normalized_landmarks, + Stream> image_size, + absl::Span + 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 +#include + +#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 landmarks = graph.In("LANDMARKS").Cast(); + Stream> tensors = + ConvertLandmarksToTensor(landmarks, + {LandmarksToTensorCalculatorOptions::X, + LandmarksToTensorCalculatorOptions::Y, + LandmarksToTensorCalculatorOptions::Z}, + /*flatten=*/true, graph); + tensors.SetName("tensors"); + + EXPECT_THAT(graph.GetConfig(), + EqualsProto(ParseTextProtoOrDie(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 landmarks = + graph.In("LANDMARKS").Cast(); + Stream> image_size = + graph.In("IMAGE_SIZE").Cast>(); + Stream> tensors = ConvertNormalizedLandmarksToTensor( + landmarks, image_size, + {LandmarksToTensorCalculatorOptions::X, + LandmarksToTensorCalculatorOptions::Y, + LandmarksToTensorCalculatorOptions::Z}, + /*flatten=*/false, graph); + tensors.SetName("tensors"); + + EXPECT_THAT(graph.GetConfig(), + EqualsProto(ParseTextProtoOrDie(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