Introduce SetJointsVisibilityCalculator

PiperOrigin-RevId: 570745171

mediapipe/calculators/util/BUILD

@@ -1659,6 +1659,49 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "set_joints_visibility_calculator",
+    srcs = ["set_joints_visibility_calculator.cc"],
+    hdrs = ["set_joints_visibility_calculator.h"],
+    deps = [
+        ":set_joints_visibility_calculator_cc_proto",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/api2:node",
+        "//mediapipe/framework/formats:body_rig_cc_proto",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/port:ret_check",
+        "//mediapipe/framework/port:status",
+    ],
+    alwayslink = 1,
+)
+
+mediapipe_proto_library(
+    name = "set_joints_visibility_calculator_proto",
+    srcs = ["set_joints_visibility_calculator.proto"],
+    deps = [
+        "//mediapipe/framework:calculator_options_proto",
+        "//mediapipe/framework:calculator_proto",
+    ],
+)
+
+cc_test(
+    name = "set_joints_visibility_calculator_test",
+    srcs = ["set_joints_visibility_calculator_test.cc"],
+    deps = [
+        ":set_joints_visibility_calculator",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:calculator_runner",
+        "//mediapipe/framework:packet",
+        "//mediapipe/framework/formats:body_rig_cc_proto",
+        "//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",
+        "@com_google_absl//absl/types:optional",
+    ],
+)
+
 cc_library(
     name = "pass_through_or_empty_detection_vector_calculator",
     srcs = ["pass_through_or_empty_detection_vector_calculator.cc"],

mediapipe/calculators/util/set_joints_visibility_calculator.cc

@@ -0,0 +1,108 @@
+// 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/set_joints_visibility_calculator.h"
+
+#include <algorithm>
+#include <optional>
+#include <utility>
+
+#include "mediapipe/calculators/util/set_joints_visibility_calculator.pb.h"
+#include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/body_rig.pb.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
+#include "mediapipe/framework/port/ret_check.h"
+
+namespace mediapipe {
+namespace api2 {
+
+namespace {}  // namespace
+
+class SetJointsVisibilityCalculatorImpl
+    : public NodeImpl<SetJointsVisibilityCalculator> {
+ public:
+  absl::Status Open(CalculatorContext* cc) override {
+    options_ = cc->Options<SetJointsVisibilityCalculatorOptions>();
+    return absl::OkStatus();
+  }
+
+  absl::Status Process(CalculatorContext* cc) override {
+    // Skip if Joints are empty.
+    if (kInJoints(cc).IsEmpty()) {
+      return absl::OkStatus();
+    }
+
+    // Get joints.
+    const JointList& in_joints = kInJoints(cc).Get();
+    RET_CHECK_EQ(in_joints.joint_size(), options_.mapping_size())
+        << "Number of joints doesn't match number of mappings";
+
+    // Get landmarks.
+    RET_CHECK(!kInLandmarks(cc).IsEmpty()) << "Landmarks must be provided";
+    const LandmarkList& in_landmarks = kInLandmarks(cc).Get();
+
+    // Set joints visibility.
+    JointList out_joints;
+    for (int i = 0; i < in_joints.joint_size(); ++i) {
+      // Initialize output joint.
+      Joint* out_joint = out_joints.add_joint();
+      *out_joint = in_joints.joint(i);
+
+      // Get visibility. But only if it exists in the source landmark(s).
+      std::optional<float> visibility;
+      auto& mapping = options_.mapping(i);
+      if (mapping.has_unchanged()) {
+        continue;
+      } else if (mapping.has_copy()) {
+        const int idx = mapping.copy().idx();
+        RET_CHECK(idx >= 0 && idx < in_landmarks.landmark_size())
+            << "Landmark index out of range";
+        if (in_landmarks.landmark(idx).has_visibility()) {
+          visibility = in_landmarks.landmark(idx).visibility();
+        }
+      } else if (mapping.has_highest()) {
+        RET_CHECK_GT(mapping.highest().idx_size(), 0) << "No indexes provided";
+        for (int idx : mapping.highest().idx()) {
+          RET_CHECK(idx >= 0 && idx < in_landmarks.landmark_size())
+              << "Landmark index out of range";
+          if (in_landmarks.landmark(idx).has_visibility()) {
+            const float landmark_visibility =
+                in_landmarks.landmark(idx).visibility();
+            visibility = visibility.has_value()
+                             ? std::max(visibility.value(), landmark_visibility)
+                             : landmark_visibility;
+          }
+        }
+      } else {
+        RET_CHECK_FAIL() << "Unknown mapping";
+      }
+
+      // Set visibility. But only if it was possible to obtain it.
+      if (visibility.has_value()) {
+        out_joint->set_visibility(visibility.value());
+      }
+    }
+
+    kOutJoints(cc).Send(std::move(out_joints));
+    return absl::OkStatus();
+  }
+
+ private:
+  SetJointsVisibilityCalculatorOptions options_;
+};
+MEDIAPIPE_NODE_IMPLEMENTATION(SetJointsVisibilityCalculatorImpl);
+
+}  // namespace api2
+}  // namespace mediapipe

mediapipe/calculators/util/set_joints_visibility_calculator.h

@@ -0,0 +1,68 @@
+// 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_SET_JOINTS_VISIBILITY_CALCULATOR_H_
+#define MEDIAPIPE_CALCULATORS_UTIL_SET_JOINTS_VISIBILITY_CALCULATOR_H_
+
+#include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/body_rig.pb.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
+
+namespace mediapipe {
+namespace api2 {
+
+// A calculator set Joints visibility from Landmarks.
+//
+// Calculator allows to either copy visibility right from the landmark or
+// somehow combine visibilities of several landmarks.
+//
+// Input:
+//   JOINTS - JointList
+//     Joints to to update visibility.
+//   LANDMARKS - LandmarkList
+//     Landmarks to take visibility from.
+//
+// Output:
+//   JOINTS - JointList
+//     Joints with updated visibility.
+//
+// Example:
+//   node {
+//     calculator: "SetJointsVisibilityCalculator"
+//     input_stream: "JOINTS:joints"
+//     input_stream: "LANDMARKS:landmarks"
+//     output_stream: "JOINTS:joints_with_visibility"
+//     options: {
+//       [mediapipe.SetJointsVisibilityCalculatorOptions.ext] {
+//         mapping: [
+//           { copy: { idx: 0 } },
+//           { highest: { idx: [5, 6] } }
+//         ]
+//       }
+//     }
+//   }
+class SetJointsVisibilityCalculator : public NodeIntf {
+ public:
+  static constexpr Input<mediapipe::JointList> kInJoints{"JOINTS"};
+  static constexpr Input<mediapipe::LandmarkList> kInLandmarks{"LANDMARKS"};
+  static constexpr Output<mediapipe::JointList> kOutJoints{"JOINTS"};
+  MEDIAPIPE_NODE_INTERFACE(SetJointsVisibilityCalculator, kInJoints,
+                           kInLandmarks, kOutJoints);
+};
+
+}  // namespace api2
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_CALCULATORS_UTIL_SET_JOINTS_VISIBILITY_CALCULATOR_H_

mediapipe/calculators/util/set_joints_visibility_calculator.proto

@@ -0,0 +1,55 @@
+// 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 SetJointsVisibilityCalculatorOptions {
+  extend CalculatorOptions {
+    optional SetJointsVisibilityCalculatorOptions ext = 406440186;
+  }
+
+  // Mapping that tells where to take visibility for the joint.
+  message Mapping {
+    // Keep visibility unchanged.
+    message Unchanged {}
+
+    // Copy visibility as is from the given landmark.
+    message Copy {
+      // Index of the landmark.
+      optional int32 idx = 1;
+    }
+
+    // Take the highest visibility among the given landmarks.
+    message Highest {
+      // Indexes of landmarks to take the highest visibility value from. At
+      // least one index must be provided.
+      repeated int32 idx = 1 [packed = true];
+    }
+
+    oneof mapping {
+      Unchanged unchanged = 1;
+      Copy copy = 2;
+      Highest highest = 3;
+    }
+  }
+
+  // Mapping that tells where to take visibility for each joint.
+  // Number of mappings must be equal to number of provided joints. Each mapping
+  // must contain exactly one rule for how to set the joint visibility.
+  repeated Mapping mapping = 1;
+}

mediapipe/calculators/util/set_joints_visibility_calculator_test.cc

@@ -0,0 +1,155 @@
+// 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 "absl/types/optional.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/calculator_runner.h"
+#include "mediapipe/framework/formats/body_rig.pb.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;
+
+struct SetJointsVisibilityTestCase {
+  std::string test_name;
+  std::string mapping;
+  std::vector<absl::optional<float>> in_joint_visibilities;
+  std::vector<absl::optional<float>> landmark_visibilities;
+  std::vector<absl::optional<float>> out_joint_visibilities;
+};
+
+using SetJointsVisibilityTest =
+    ::testing::TestWithParam<SetJointsVisibilityTestCase>;
+
+TEST_P(SetJointsVisibilityTest, SetJointsVisibilityTest) {
+  const SetJointsVisibilityTestCase& tc = GetParam();
+
+  // Prepare graph.
+  mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(absl::Substitute(
+      R"(
+      calculator: "SetJointsVisibilityCalculator"
+      input_stream: "JOINTS:joints"
+      input_stream: "LANDMARKS:landmarks"
+      output_stream: "JOINTS:joints_with_visibility"
+      options: {
+        [mediapipe.SetJointsVisibilityCalculatorOptions.ext] {
+          mapping: [
+            $0
+          ]
+        }
+      }
+  )",
+      tc.mapping)));
+
+  // Prepare joints.
+  JointList in_joints;
+  for (auto vis_opt : tc.in_joint_visibilities) {
+    Joint* joint = in_joints.add_joint();
+    if (vis_opt) {
+      joint->set_visibility(vis_opt.value());
+    }
+  }
+
+  // Prepare landmarks.
+  LandmarkList landmarks;
+  for (auto vis_opt : tc.landmark_visibilities) {
+    Landmark* lmk = landmarks.add_landmark();
+    if (vis_opt) {
+      lmk->set_visibility(vis_opt.value());
+    }
+  }
+
+  // Send joints and landmarks to the graph.
+  runner.MutableInputs()->Tag("JOINTS").packets.push_back(
+      MakePacket<JointList>(std::move(in_joints)).At(mediapipe::Timestamp(0)));
+  runner.MutableInputs()
+      ->Tag("LANDMARKS")
+      .packets.push_back(MakePacket<LandmarkList>(std::move(landmarks))
+                             .At(mediapipe::Timestamp(0)));
+
+  // Run the graph.
+  MP_ASSERT_OK(runner.Run());
+
+  const auto& output_packets = runner.Outputs().Tag("JOINTS").packets;
+  EXPECT_EQ(1, output_packets.size());
+
+  const auto& out_joints = output_packets[0].Get<JointList>();
+  EXPECT_EQ(out_joints.joint_size(), tc.out_joint_visibilities.size());
+  for (int i = 0; i < out_joints.joint_size(); ++i) {
+    const Joint& joint = out_joints.joint(i);
+    auto expected_vis_opt = tc.out_joint_visibilities[i];
+    if (expected_vis_opt) {
+      EXPECT_NEAR(joint.visibility(), expected_vis_opt.value(), 1e-5);
+    } else {
+      EXPECT_FALSE(joint.has_visibility());
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    SetJointsVisibilityTests, SetJointsVisibilityTest,
+    testing::ValuesIn<SetJointsVisibilityTestCase>({
+        {"Empty_LandmarksAndJoints", "", {}, {}, {}},
+        {"Empty_Joints", "", {}, {1, 2, 3}, {}},
+        {"Empty_Landmarks",
+         "{ unchanged: {} }, { unchanged: {} }, { unchanged: {} }",
+         {1, 2, absl::nullopt},
+         {},
+         {1, 2, absl::nullopt}},
+
+        {"Mapping_Unchanged", "{ unchanged: {} }", {1}, {2}, {1}},
+        {"Mapping_Unchanged_UnsetJointVisRemainsUnset",
+         "{ unchanged: {} }",
+         {absl::nullopt},
+         {2},
+         {absl::nullopt}},
+
+        {"Mapping_Copy", "{ copy: { idx: 0 } }", {1}, {2}, {2}},
+        {"Mapping_Copy_UnsetLmkVisResultsIntoZeroJointVis",
+         "{ copy: { idx: 0 } }",
+         {absl::nullopt},
+         {absl::nullopt},
+         {0}},
+
+        {"Mapping_Highest",
+         "{ highest: { idx: [0, 1, 2] } }",
+         {absl::nullopt},
+         {2, 4, 3},
+         {4}},
+        {"Mapping_Highest_UnsetLmkIsIgnored",
+         "{ highest: { idx: [0, 1, 2] } }",
+         {absl::nullopt},
+         {-2, absl::nullopt, -3},
+         {-2}},
+    }),
+    [](const testing::TestParamInfo<SetJointsVisibilityTest::ParamType>& info) {
+      return info.param.test_name;
+    });
+
+}  // namespace
+}  // namespace api2
+}  // namespace mediapipe

mediapipe/framework/tool/mediapipe_proto_allowlist.bzl

@@ -248,7 +248,6 @@ rewrite_target_list = [
     "segmenter_proto",
     "sequence_shift_calculator_proto",
     "set_alpha_calculator_proto",
-    "set_joints_visibility_calculator_proto",
     "sharpen_calculator_gl_proto",
     "simple_calculator_proto",
     "single_shot_detector_gpu_cpu_proto",