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(
|
cc_library(
|
||||||
name = "landmarks_projection",
|
name = "landmarks_projection",
|
||||||
srcs = ["landmarks_projection.cc"],
|
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