From 9edb4cd75353eed8ffd180437bee627ba13d08d1 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Wed, 27 Sep 2023 21:09:59 -0700 Subject: [PATCH] Introduce LandmarksTransformationCalculator PiperOrigin-RevId: 569050686 --- mediapipe/calculators/util/BUILD | 41 +++++ .../landmarks_transformation_calculator.cc | 142 +++++++++++++++ .../landmarks_transformation_calculator.h | 62 +++++++ .../landmarks_transformation_calculator.proto | 45 +++++ ...andmarks_transformation_calculator_test.cc | 167 ++++++++++++++++++ 5 files changed, 457 insertions(+) create mode 100644 mediapipe/calculators/util/landmarks_transformation_calculator.cc create mode 100644 mediapipe/calculators/util/landmarks_transformation_calculator.h create mode 100644 mediapipe/calculators/util/landmarks_transformation_calculator.proto create mode 100644 mediapipe/calculators/util/landmarks_transformation_calculator_test.cc diff --git a/mediapipe/calculators/util/BUILD b/mediapipe/calculators/util/BUILD index 328b24fd..4186cbea 100644 --- a/mediapipe/calculators/util/BUILD +++ b/mediapipe/calculators/util/BUILD @@ -1575,3 +1575,44 @@ cc_test( "@com_google_absl//absl/status", ], ) + +mediapipe_proto_library( + name = "landmarks_transformation_calculator_proto", + srcs = ["landmarks_transformation_calculator.proto"], + deps = [ + "//mediapipe/framework:calculator_options_proto", + "//mediapipe/framework:calculator_proto", + ], +) + +cc_library( + name = "landmarks_transformation_calculator", + srcs = ["landmarks_transformation_calculator.cc"], + hdrs = ["landmarks_transformation_calculator.h"], + deps = [ + ":landmarks_transformation_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/api2:node", + "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/port:ret_check", + "//mediapipe/framework/port:statusor", + ], + alwayslink = 1, +) + +cc_test( + name = "landmarks_transformation_calculator_test", + srcs = ["landmarks_transformation_calculator_test.cc"], + deps = [ + ":landmarks_transformation_calculator", + ":landmarks_transformation_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:calculator_runner", + "//mediapipe/framework:packet", + "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/port:status_matchers", + "@com_google_absl//absl/strings", + ], +) diff --git a/mediapipe/calculators/util/landmarks_transformation_calculator.cc b/mediapipe/calculators/util/landmarks_transformation_calculator.cc new file mode 100644 index 00000000..96eaba54 --- /dev/null +++ b/mediapipe/calculators/util/landmarks_transformation_calculator.cc @@ -0,0 +1,142 @@ +// Copyright 2023 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. + +#include "mediapipe/calculators/util/landmarks_transformation_calculator.h" + +#include + +#include "mediapipe/calculators/util/landmarks_transformation_calculator.pb.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/port/ret_check.h" +#include "mediapipe/framework/port/statusor.h" + +namespace mediapipe { +namespace api2 { + +namespace { + +StatusOr NormalizeTranslation(const LandmarkList& in_landmarks) { + RET_CHECK_GT(in_landmarks.landmark_size(), 0); + + double x_sum = 0.0f; + double y_sum = 0.0f; + double z_sum = 0.0f; + for (auto& in_landmark : in_landmarks.landmark()) { + x_sum += in_landmark.x(); + y_sum += in_landmark.y(); + z_sum += in_landmark.z(); + } + + float x_mean = x_sum / in_landmarks.landmark_size(); + float y_mean = y_sum / in_landmarks.landmark_size(); + float z_mean = z_sum / in_landmarks.landmark_size(); + + LandmarkList out_landmarks; + for (auto& in_landmark : in_landmarks.landmark()) { + auto* out_landmark = out_landmarks.add_landmark(); + *out_landmark = in_landmark; + out_landmark->set_x(in_landmark.x() - x_mean); + out_landmark->set_y(in_landmark.y() - y_mean); + out_landmark->set_z(in_landmark.z() - z_mean); + } + + return out_landmarks; +} + +StatusOr FlipAxis( + const LandmarkList& in_landmarks, + const LandmarksTransformationCalculatorOptions::FlipAxis& options) { + float x_mul = options.flip_x() ? -1 : 1; + float y_mul = options.flip_y() ? -1 : 1; + float z_mul = options.flip_z() ? -1 : 1; + + LandmarkList out_landmarks; + for (auto& in_landmark : in_landmarks.landmark()) { + auto* out_landmark = out_landmarks.add_landmark(); + *out_landmark = in_landmark; + out_landmark->set_x(in_landmark.x() * x_mul); + out_landmark->set_y(in_landmark.y() * y_mul); + out_landmark->set_z(in_landmark.z() * z_mul); + } + + return out_landmarks; +} + +} // namespace + +class LandmarksTransformationCalculatorImpl + : public NodeImpl { + public: + static absl::Status UpdateContract(CalculatorContract* cc) { + // Check that if options input stream is connected there should be no static + // options in calculator. Currently there is no such functionality, so we'll + // just check for the number of transforms. + if (kInOptions(cc).IsConnected()) { + RET_CHECK_EQ(cc->Options() + .transformation_size(), + 0); + } + return absl::OkStatus(); + } + + absl::Status Open(CalculatorContext* cc) override { + options_ = cc->Options(); + return absl::OkStatus(); + } + + absl::Status Process(CalculatorContext* cc) override { + if (kInLandmarks(cc).IsEmpty()) { + return absl::OkStatus(); + } + + // Get transformation options for either calculator parameters or input + // stream. Input stream has higher priority. + LandmarksTransformationCalculatorOptions options; + if (kInOptions(cc).IsConnected()) { + // If input stream is connected but is empty - use no transformations and + // return landmarks as is. + if (!kInOptions(cc).IsEmpty()) { + options = kInOptions(cc).Get(); + } + } else { + options = options_; + } + + LandmarkList landmarks = kInLandmarks(cc).Get(); + + for (auto& transformation : options.transformation()) { + if (transformation.has_normalize_translation()) { + ASSIGN_OR_RETURN(landmarks, NormalizeTranslation(landmarks)); + } else if (transformation.has_flip_axis()) { + ASSIGN_OR_RETURN(landmarks, + FlipAxis(landmarks, transformation.flip_axis())); + } else { + RET_CHECK_FAIL() << "Unknown landmarks transformation"; + } + } + + kOutLandmarks(cc).Send(std::move(landmarks)); + + return absl::OkStatus(); + } + + private: + LandmarksTransformationCalculatorOptions options_; +}; +MEDIAPIPE_NODE_IMPLEMENTATION(LandmarksTransformationCalculatorImpl); + +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/calculators/util/landmarks_transformation_calculator.h b/mediapipe/calculators/util/landmarks_transformation_calculator.h new file mode 100644 index 00000000..5bc31492 --- /dev/null +++ b/mediapipe/calculators/util/landmarks_transformation_calculator.h @@ -0,0 +1,62 @@ +// Copyright 2023 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. + +#ifndef MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TRANSFORMATION_CALCULATOR_H_ +#define MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TRANSFORMATION_CALCULATOR_H_ + +#include "mediapipe/calculators/util/landmarks_transformation_calculator.pb.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/landmark.pb.h" + +namespace mediapipe { +namespace api2 { + +// A calculator to transform landmarks. +// +// Input: +// LANDMARKS - LandmarkList +// Landmarks to transform. +// +// Output: +// LANDMARKS - LandmarkList +// Transformed landmarks. +// +// Example: +// node { +// calculator: "LandmarksTransformationCalculator" +// input_stream: "LANDMARKS:in_landmarks" +// output_stream: "LANDMARKS:out_landmarks" +// options: { +// [mediapipe.LandmarksTransformationCalculatorOptions.ext] { +// transformation: { normalize_translation: {} } +// transformation: { flip_axis: { flip_x: true } } +// } +// } +// } +class LandmarksTransformationCalculator : public NodeIntf { + public: + static constexpr Input kInLandmarks{"LANDMARKS"}; + static constexpr Input< + mediapipe::LandmarksTransformationCalculatorOptions>::Optional kInOptions{ + "OPTIONS"}; + static constexpr Output kOutLandmarks{"LANDMARKS"}; + MEDIAPIPE_NODE_INTERFACE(LandmarksTransformationCalculator, kInLandmarks, + kInOptions, kOutLandmarks); +}; + +} // namespace api2 +} // namespace mediapipe + +#endif // MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TRANSFORMATION_CALCULATOR_H_ diff --git a/mediapipe/calculators/util/landmarks_transformation_calculator.proto b/mediapipe/calculators/util/landmarks_transformation_calculator.proto new file mode 100644 index 00000000..0ced6871 --- /dev/null +++ b/mediapipe/calculators/util/landmarks_transformation_calculator.proto @@ -0,0 +1,45 @@ +// Copyright 2023 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. + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; + +message LandmarksTransformationCalculatorOptions { + extend CalculatorOptions { + optional LandmarksTransformationCalculatorOptions ext = 421309928; + } + + // Normalize translation so that center of mass is in (0, 0, 0). + message NormalizeTranslation {} + + // Flip axis by multiplying coordinates along it by `-1`. + message FlipAxis { + optional bool flip_x = 1 [default = false]; + optional bool flip_y = 2 [default = false]; + optional bool flip_z = 3 [default = false]; + } + + message Transformation { + oneof transformation { + NormalizeTranslation normalize_translation = 1; + FlipAxis flip_axis = 2; + } + } + + // Transformations applied in given order. + repeated Transformation transformation = 1; +} diff --git a/mediapipe/calculators/util/landmarks_transformation_calculator_test.cc b/mediapipe/calculators/util/landmarks_transformation_calculator_test.cc new file mode 100644 index 00000000..76db742c --- /dev/null +++ b/mediapipe/calculators/util/landmarks_transformation_calculator_test.cc @@ -0,0 +1,167 @@ +// Copyright 2023 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. + +#include +#include +#include + +#include "absl/strings/substitute.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/calculator_runner.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/packet.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 { +namespace api2 { +namespace { + +using Node = ::mediapipe::CalculatorGraphConfig::Node; + +Landmark CreateLandmark(float x, float y, float z) { + Landmark lmk; + lmk.set_x(x); + lmk.set_y(y); + lmk.set_z(z); + return lmk; +} + +Landmark CreateLandmark(float x, float y, float z, float visibility, + float presence) { + Landmark lmk; + lmk.set_x(x); + lmk.set_y(y); + lmk.set_z(z); + lmk.set_visibility(visibility); + lmk.set_presence(presence); + return lmk; +} + +struct LandmarksTransformationestCase { + std::string test_name; + std::string transformations; + std::vector in_landmarks; + std::vector out_landmarks; +}; + +using LandmarksTransformationest = + ::testing::TestWithParam; + +TEST_P(LandmarksTransformationest, LandmarksTransformationest) { + const LandmarksTransformationestCase& tc = GetParam(); + + // Prepare graph. + mediapipe::CalculatorRunner runner(ParseTextProtoOrDie(absl::Substitute( + R"( + calculator: "LandmarksTransformationCalculator" + input_stream: "LANDMARKS:in_landmarks" + output_stream: "LANDMARKS:out_landmarks" + options: { + [mediapipe.LandmarksTransformationCalculatorOptions.ext] { + $0 + } + } + )", + tc.transformations))); + + // In landmarks. + LandmarkList in_landmarks; + for (auto& lmk : tc.in_landmarks) { + *in_landmarks.add_landmark() = lmk; + } + + // Send landmarks to the graph. + runner.MutableInputs() + ->Tag("LANDMARKS") + .packets.push_back(MakePacket(std::move(in_landmarks)) + .At(mediapipe::Timestamp(0))); + + // Run the graph. + MP_ASSERT_OK(runner.Run()); + + const auto& output_packets = runner.Outputs().Tag("LANDMARKS").packets; + EXPECT_EQ(1, output_packets.size()); + + const auto& out_landmarks = output_packets[0].Get(); + EXPECT_EQ(out_landmarks.landmark_size(), tc.out_landmarks.size()); + for (int i = 0; i < out_landmarks.landmark_size(); ++i) { + auto& lmk = out_landmarks.landmark(i); + auto& exp_lmk = tc.out_landmarks[i]; + EXPECT_EQ(lmk.x(), exp_lmk.x()) << "Unexpected lmk[" << i << "].x"; + EXPECT_EQ(lmk.y(), exp_lmk.y()) << "Unexpected lmk[" << i << "].y"; + EXPECT_EQ(lmk.z(), exp_lmk.z()) << "Unexpected lmk[" << i << "].z"; + if (exp_lmk.has_visibility()) { + EXPECT_EQ(lmk.visibility(), exp_lmk.visibility()) + << "Unexpected lmk[" << i << "].visibility"; + } + if (exp_lmk.has_presence()) { + EXPECT_EQ(lmk.presence(), exp_lmk.presence()) + << "Unexpected lmk[" << i << "].presence"; + } + } +} + +INSTANTIATE_TEST_SUITE_P( + LandmarksTransformationests, LandmarksTransformationest, + testing::ValuesIn({ + {"NoTransformations", + "", + {CreateLandmark(1, 2, 3), CreateLandmark(4, 5, 6)}, + {CreateLandmark(1, 2, 3), CreateLandmark(4, 5, 6)}}, + + {"NormalizeTranslation_OneLandmark", + "transformation: { normalize_translation: {} }", + {CreateLandmark(2, 2, 2)}, + {CreateLandmark(0, 0, 0)}}, + {"NormalizeTranslation_TwoLandmarks", + "transformation: { normalize_translation: {} }", + {CreateLandmark(2, 2, 2), CreateLandmark(4, 4, 4)}, + {CreateLandmark(-1, -1, -1), CreateLandmark(1, 1, 1)}}, + {"NormalizeTranslation_ThreeLandmarks", + "transformation: { normalize_translation: {} }", + {CreateLandmark(2, 2, 2), CreateLandmark(4, 4, 4), + CreateLandmark(9, 9, 9)}, + {CreateLandmark(-3, -3, -3), CreateLandmark(-1, -1, -1), + CreateLandmark(4, 4, 4)}}, + {"NormalizeTranslation_VisibilityAndPresence", + "transformation: { normalize_translation: {} }", + {CreateLandmark(0, 0, 0, 4, 5)}, + {CreateLandmark(0, 0, 0, 4, 5)}}, + + {"FlipAxis_X", + "transformation: { flip_axis: { flip_x: true } }", + {CreateLandmark(2, 2, 2)}, + {CreateLandmark(-2, 2, 2)}}, + {"FlipAxis_Y", + "transformation: { flip_axis: { flip_y: true } }", + {CreateLandmark(2, 2, 2)}, + {CreateLandmark(2, -2, 2)}}, + {"FlipAxis_Z", + "transformation: { flip_axis: { flip_z: true } }", + {CreateLandmark(2, 2, 2)}, + {CreateLandmark(2, 2, -2)}}, + {"FlipAxis_VisibilityAndPresence", + "transformation: { flip_axis: { flip_x: true } }", + {CreateLandmark(0, 0, 0, 4, 5)}, + {CreateLandmark(0, 0, 0, 4, 5)}}, + }), + [](const testing::TestParamInfo& + info) { return info.param.test_name; }); + +} // namespace +} // namespace api2 +} // namespace mediapipe