Internal change

PiperOrigin-RevId: 477538515

mediapipe/calculators/core/BUILD

@@ -1294,8 +1294,8 @@ cc_library(
     deps = [
         ":get_vector_item_calculator_cc_proto",
         "//mediapipe/framework:calculator_framework",
-        "//mediapipe/framework:packet",
         "//mediapipe/framework/api2:node",
+        "//mediapipe/framework/api2:packet",
         "//mediapipe/framework/api2:port",
         "//mediapipe/framework/formats:classification_cc_proto",
         "//mediapipe/framework/formats:landmark_cc_proto",
@@ -1319,6 +1319,32 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "vector_indices_calculator",
+    srcs = ["vector_indices_calculator.cc"],
+    hdrs = ["vector_indices_calculator.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/api2:node",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/port:status",
+    ],
+    alwayslink = 1,
+)
+
+cc_test(
+    name = "vector_indices_calculator_test",
+    srcs = ["vector_indices_calculator_test.cc"],
+    deps = [
+        ":vector_indices_calculator",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:calculator_runner",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/port:parse_text_proto",
+    ],
+)
+
 cc_library(
     name = "vector_size_calculator",
     srcs = ["vector_size_calculator.cc"],

mediapipe/calculators/core/end_loop_calculator.cc

@@ -40,6 +40,9 @@ REGISTER_CALCULATOR(EndLoopNormalizedLandmarkListVectorCalculator);
 typedef EndLoopCalculator<std::vector<bool>> EndLoopBooleanCalculator;
 REGISTER_CALCULATOR(EndLoopBooleanCalculator);
 
+typedef EndLoopCalculator<std::vector<float>> EndLoopFloatCalculator;
+REGISTER_CALCULATOR(EndLoopFloatCalculator);
+
 typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
     EndLoopRenderDataCalculator;
 REGISTER_CALCULATOR(EndLoopRenderDataCalculator);

mediapipe/calculators/core/get_vector_item_calculator.cc

@@ -24,6 +24,10 @@ using GetLandmarkListVectorItemCalculator =
     GetVectorItemCalculator<mediapipe::LandmarkList>;
 REGISTER_CALCULATOR(GetLandmarkListVectorItemCalculator);
 
+using GetNormalizedLandmarkListVectorItemCalculator =
+    GetVectorItemCalculator<mediapipe::NormalizedLandmarkList>;
+REGISTER_CALCULATOR(GetNormalizedLandmarkListVectorItemCalculator);
+
 using GetClassificationListVectorItemCalculator =
     GetVectorItemCalculator<mediapipe::ClassificationList>;
 REGISTER_CALCULATOR(GetClassificationListVectorItemCalculator);

mediapipe/calculators/core/get_vector_item_calculator.h

@@ -19,6 +19,7 @@
 
 #include "mediapipe/calculators/core/get_vector_item_calculator.pb.h"
 #include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/api2/packet.h"
 #include "mediapipe/framework/api2/port.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/port/ret_check.h"
@@ -58,7 +59,7 @@ template <typename T>
 class GetVectorItemCalculator : public Node {
  public:
   static constexpr Input<std::vector<T>> kIn{"VECTOR"};
-  static constexpr Input<int>::Optional kIdx{"INDEX"};
+  static constexpr Input<OneOf<int, uint64_t>>::Optional kIdx{"INDEX"};
   static constexpr Output<T> kOut{"ITEM"};
 
   MEDIAPIPE_NODE_CONTRACT(kIn, kIdx, kOut);
@@ -80,7 +81,9 @@ class GetVectorItemCalculator : public Node {
 
     int idx = 0;
     if (kIdx(cc).IsConnected() && !kIdx(cc).IsEmpty()) {
-      idx = kIdx(cc).Get();
+      idx = kIdx(cc).Visit(
+          [](uint64_t idx_uint64_t) { return static_cast<int>(idx_uint64_t); },
+          [](int idx_int) { return idx_int; });
     } else if (options.has_item_index()) {
       idx = options.item_index();
     } else {

mediapipe/calculators/core/get_vector_item_calculator_test.cc

@@ -227,4 +227,15 @@ TEST(TestGetIntVectorItemCalculatorTest, IndexOptionsTwoTimestamps) {
               testing::ElementsAre(TimestampValue(1), TimestampValue(2)));
 }
 
+TEST(TestGetIntVectorItemCalculatorTest, IndexUint64) {
+  CalculatorRunner runner = MakeRunnerWithStream();
+  const std::vector<int> inputs = {1, 2, 3};
+  const uint64_t index = 1;
+  AddInputVector(runner, inputs, 1);
+  AddInputIndex(runner, index, 1);
+  MP_ASSERT_OK(runner.Run());
+  const std::vector<Packet>& outputs = runner.Outputs().Tag("ITEM").packets;
+  EXPECT_THAT(outputs, testing::ElementsAre(IntPacket(inputs[index])));
+}
+
 }  // namespace mediapipe

mediapipe/calculators/core/vector_indices_calculator.cc

@@ -0,0 +1,33 @@
+// Copyright 2022 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/core/vector_indices_calculator.h"
+
+#include "mediapipe/framework/formats/landmark.pb.h"
+
+namespace mediapipe {
+namespace api2 {
+
+using IntVectorIndicesCalculator = VectorIndicesCalculator<int>;
+REGISTER_CALCULATOR(IntVectorIndicesCalculator);
+
+using Uint64tVectorIndicesCalculator = VectorIndicesCalculator<uint64_t>;
+REGISTER_CALCULATOR(Uint64tVectorIndicesCalculator);
+
+using NormalizedLandmarkListVectorIndicesCalculator =
+    VectorIndicesCalculator<mediapipe::NormalizedLandmarkList>;
+REGISTER_CALCULATOR(NormalizedLandmarkListVectorIndicesCalculator);
+
+}  // namespace api2
+}  // namespace mediapipe

mediapipe/calculators/core/vector_indices_calculator.h

@@ -0,0 +1,65 @@
+// Copyright 2022 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_CORE_VECTOR_INDICES_CALCULATOR_H_
+#define MEDIAPIPE_CALCULATORS_CORE_VECTOR_INDICES_CALCULATOR_H_
+
+#include <optional>
+
+#include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/port/status.h"
+
+namespace mediapipe {
+namespace api2 {
+// Calculator that takes a vector and consturct an index range vector based on
+// the size of the input vector.
+//
+// Inputs:
+//   VECTOR - std::vector<T>
+//     Vector whose range of indices to return.
+//
+// Outputs:
+//   INDICES - std::vector<int>
+//     Indices vector of the input vector.
+//
+// Example config:
+//  node {
+//    calculator: "{SpecificType}VectorIndicesCalculator"
+//    input_stream: "VECTOR:vector"
+//    output_stream: "INDICES:indices"
+//  }
+//
+template <typename T>
+class VectorIndicesCalculator : public Node {
+ public:
+  static constexpr Input<std::vector<T>> kVector{"VECTOR"};
+  static constexpr Output<std::vector<int>> kRange{"INDICES"};
+
+  MEDIAPIPE_NODE_CONTRACT(kVector, kRange);
+
+  absl::Status Process(CalculatorContext* cc) final {
+    // Get the size of the input vector.
+    const int vector_size = kVector(cc).Get().size();
+    std::vector<int> out_idxs(vector_size);
+    std::iota(out_idxs.begin(), out_idxs.end(), 0);
+    kRange(cc).Send(out_idxs);
+    return absl::OkStatus();
+  }
+};
+
+}  // namespace api2
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_CALCULATORS_CORE_VECTOR_INDICES_CALCULATOR_H_

mediapipe/calculators/core/vector_indices_calculator_test.cc

@@ -0,0 +1,87 @@
+// Copyright 2022 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/core/vector_indices_calculator.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/calculator_runner.h"
+#include "mediapipe/framework/port/gmock.h"
+#include "mediapipe/framework/port/status_matchers.h"
+
+namespace mediapipe {
+
+namespace {
+
+using ::testing::TestParamInfo;
+using ::testing::TestWithParam;
+using ::testing::Values;
+
+template <typename T>
+void AddInputVector(CalculatorRunner& runner, const std::vector<T>& inputs,
+                    int timestamp) {
+  runner.MutableInputs()->Tag("VECTOR").packets.push_back(
+      MakePacket<std::vector<T>>(inputs).At(Timestamp(timestamp)));
+}
+
+template <typename T>
+struct TestParams {
+  const std::string test_name;
+  const std::vector<T> inputs;
+  const int timestamp;
+  const std::vector<int> expected_indices;
+};
+
+class IntVectorIndicesCalculatorTest
+    : public testing::TestWithParam<TestParams<int>> {};
+
+TEST_P(IntVectorIndicesCalculatorTest, Succeeds) {
+  CalculatorRunner runner = CalculatorRunner(R"(
+    calculator: "IntVectorIndicesCalculator"
+    input_stream: "VECTOR:vector_stream"
+    output_stream: "INDICES:indices_stream"
+  )");
+  const std::vector<int>& inputs = GetParam().inputs;
+  std::vector<int> expected_indices(inputs.size());
+  AddInputVector(runner, inputs, GetParam().timestamp);
+  MP_ASSERT_OK(runner.Run());
+  const std::vector<Packet>& outputs = runner.Outputs().Tag("INDICES").packets;
+  EXPECT_EQ(1, outputs.size());
+  EXPECT_THAT(outputs[0].Get<std::vector<int>>(),
+              testing::ElementsAreArray(GetParam().expected_indices));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    IntVectorIndicesCalculatorTest, IntVectorIndicesCalculatorTest,
+    Values(TestParams<int>{
+               /* test_name= */ "IntVectorIndices",
+               /* inputs= */ {1, 2, 3},
+               /* timestamp= */ 1,
+               /* expected_indices= */ {0, 1, 2},
+           },
+           TestParams<int>{
+               /* test_name= */ "EmptyVector",
+               /* inputs= */ {},
+               /* timestamp= */ 1,
+               /* expected_indices= */ {},
+           }),
+    [](const TestParamInfo<IntVectorIndicesCalculatorTest::ParamType>& info) {
+      return info.param.test_name;
+    });
+
+}  // namespace
+}  // namespace mediapipe