landmarks_to_tensor stream utility function.
PiperOrigin-RevId: 569003241
This commit is contained in:
parent
8837b49026
commit
da02052c70
|
@ -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"],
|
||||
|
|
64
mediapipe/framework/api2/stream/landmarks_to_tensor.cc
Normal file
64
mediapipe/framework/api2/stream/landmarks_to_tensor.cc
Normal file
|
@ -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
|
37
mediapipe/framework/api2/stream/landmarks_to_tensor.h
Normal file
37
mediapipe/framework/api2/stream/landmarks_to_tensor.h
Normal file
|
@ -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_
|
89
mediapipe/framework/api2/stream/landmarks_to_tensor_test.cc
Normal file
89
mediapipe/framework/api2/stream/landmarks_to_tensor_test.cc
Normal file
|
@ -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
|
Loading…
Reference in New Issue
Block a user