split stream utility function.

PiperOrigin-RevId: 566722901
google-ai-edge · Sep 19, 2023 · bbf40cb · bbf40cb
1 parent 58bb2d1
commit bbf40cb
Show file tree

Hide file tree

Showing 6 changed files with 875 additions and 0 deletions.
diff --git a/mediapipe/calculators/core/BUILD b/mediapipe/calculators/core/BUILD
@@ -944,6 +944,7 @@ cc_library(
     deps = [
         ":split_vector_calculator_cc_proto",
         "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/formats:body_rig_cc_proto",
         "//mediapipe/framework/formats:landmark_cc_proto",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:status",

diff --git a/mediapipe/calculators/core/split_proto_list_calculator.cc b/mediapipe/calculators/core/split_proto_list_calculator.cc
@@ -17,6 +17,7 @@
 
 #include "mediapipe/calculators/core/split_vector_calculator.pb.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/canonical_errors.h"
 #include "mediapipe/framework/port/ret_check.h"
@@ -196,6 +197,18 @@ class SplitLandmarkListCalculator
 };
 REGISTER_CALCULATOR(SplitLandmarkListCalculator);
 
+class SplitJointListCalculator : public SplitListsCalculator<Joint, JointList> {
+ protected:
+  int ListSize(const JointList& list) const override {
+    return list.joint_size();
+  }
+  const Joint GetItem(const JointList& list, int idx) const override {
+    return list.joint(idx);
+  }
+  Joint* AddItem(JointList& list) const override { return list.add_joint(); }
+};
+REGISTER_CALCULATOR(SplitJointListCalculator);
+
 }  // namespace mediapipe
 
 // NOLINTNEXTLINE

diff --git a/mediapipe/framework/api2/stream/BUILD b/mediapipe/framework/api2/stream/BUILD
@@ -192,3 +192,40 @@ cc_test(
         "//mediapipe/framework/port:status_matchers",
     ],
 )
+
+cc_library(
+    name = "split",
+    hdrs = ["split.h"],
+    deps = [
+        "//mediapipe/calculators/core:split_proto_list_calculator",
+        "//mediapipe/calculators/core:split_vector_calculator",
+        "//mediapipe/calculators/core:split_vector_calculator_cc_proto",
+        "//mediapipe/framework/api2:builder",
+        "//mediapipe/framework/api2:port",
+        "//mediapipe/framework/formats:body_rig_cc_proto",
+        "//mediapipe/framework/formats:classification_cc_proto",
+        "//mediapipe/framework/formats:detection_cc_proto",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/formats:matrix",
+        "//mediapipe/framework/formats:rect_cc_proto",
+        "//mediapipe/framework/formats:tensor",
+        "@org_tensorflow//tensorflow/lite/c:common",
+    ],
+)
+
+cc_test(
+    name = "split_test",
+    srcs = ["split_test.cc"],
+    deps = [
+        ":split",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/api2:builder",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/formats:rect_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",
+    ],
+)