Introduce LandmarksTransformationCalculator

PiperOrigin-RevId: 569050686

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",
+    ],
+)

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 <utility>
+
+#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<LandmarkList> 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<LandmarkList> 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<LandmarksTransformationCalculator> {
+ 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<LandmarksTransformationCalculatorOptions>()
+                       .transformation_size(),
+                   0);
+    }
+    return absl::OkStatus();
+  }
+
+  absl::Status Open(CalculatorContext* cc) override {
+    options_ = cc->Options<LandmarksTransformationCalculatorOptions>();
+    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

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<mediapipe::LandmarkList> kInLandmarks{"LANDMARKS"};
+  static constexpr Input<
+      mediapipe::LandmarksTransformationCalculatorOptions>::Optional kInOptions{
+      "OPTIONS"};
+  static constexpr Output<mediapipe::LandmarkList> kOutLandmarks{"LANDMARKS"};
+  MEDIAPIPE_NODE_INTERFACE(LandmarksTransformationCalculator, kInLandmarks,
+                           kInOptions, kOutLandmarks);
+};
+
+}  // namespace api2
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TRANSFORMATION_CALCULATOR_H_

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;
+}

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 <string>
+#include <utility>
+#include <vector>
+
+#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<Landmark> in_landmarks;
+  std::vector<Landmark> out_landmarks;
+};
+
+using LandmarksTransformationest =
+    ::testing::TestWithParam<LandmarksTransformationestCase>;
+
+TEST_P(LandmarksTransformationest, LandmarksTransformationest) {
+  const LandmarksTransformationestCase& tc = GetParam();
+
+  // Prepare graph.
+  mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(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<LandmarkList>(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<LandmarkList>();
+  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<LandmarksTransformationestCase>({
+        {"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<LandmarksTransformationest::ParamType>&
+           info) { return info.param.test_name; });
+
+}  // namespace
+}  // namespace api2
+}  // namespace mediapipe