Internal change

PiperOrigin-RevId: 477538515
2022-09-28 21:32:36 +00:00
parent 6cdc6443b6
commit f8af41b1eb
236 changed files with 12865 additions and 1923 deletions
@@ -21,7 +21,9 @@ cc_library(
        ":port",
        "//mediapipe/framework:calculator_base",
        "//mediapipe/framework:calculator_contract",
+        "@com_google_absl//absl/container:btree",
        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings",
    ],
 )

@@ -4,7 +4,9 @@
 #include <string>
 #include <type_traits>

+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/strings/string_view.h"
 #include "mediapipe/framework/api2/const_str.h"
 #include "mediapipe/framework/api2/contract.h"
 #include "mediapipe/framework/api2/node.h"
@@ -46,7 +48,7 @@ struct TagIndexLocation {
 template <typename T>
 class TagIndexMap {
 public:
-  std::vector<std::unique_ptr<T>>& operator[](const std::string& tag) {
+  std::vector<std::unique_ptr<T>>& operator[](absl::string_view tag) {
    return map_[tag];
  }

@@ -72,7 +74,7 @@ class TagIndexMap {

  // Note: entries are held by a unique_ptr to ensure pointers remain valid.
  // Should use absl::flat_hash_map but ordering keys for now.
-  std::map<std::string, std::vector<std::unique_ptr<T>>> map_;
+  absl::btree_map<std::string, std::vector<std::unique_ptr<T>>> map_;
 };

 class Graph;
@@ -169,6 +171,16 @@ class SourceImpl {
    return AddTarget(dest);
  }

+  template <typename U>
+  struct AllowCast
+      : public std::integral_constant<bool, std::is_same_v<T, AnyType> &&
+                                                !std::is_same_v<T, U>> {};
+
+  template <typename U, std::enable_if_t<AllowCast<U>{}, int> = 0>
+  SourceImpl<IsSide, U> Cast() {
+    return SourceImpl<IsSide, U>(base_);
+  }
+
 private:
  // Never null.
  SourceBase* base_;
@@ -212,19 +224,19 @@ class NodeBase {
  // of its entries by index. However, for nodes without visible contracts we
  // can't know whether a tag is indexable or not, so we would need the
  // multi-port to also be usable as a port directly (representing index 0).
-  MultiSource<> Out(const std::string& tag) {
+  MultiSource<> Out(absl::string_view tag) {
    return MultiSource<>(&out_streams_[tag]);
  }

-  MultiDestination<> In(const std::string& tag) {
+  MultiDestination<> In(absl::string_view tag) {
    return MultiDestination<>(&in_streams_[tag]);
  }

-  MultiSideSource<> SideOut(const std::string& tag) {
+  MultiSideSource<> SideOut(absl::string_view tag) {
    return MultiSideSource<>(&out_sides_[tag]);
  }

-  MultiSideDestination<> SideIn(const std::string& tag) {
+  MultiSideDestination<> SideIn(absl::string_view tag) {
    return MultiSideDestination<>(&in_sides_[tag]);
  }

@@ -359,11 +371,11 @@ class PacketGenerator {
 public:
  PacketGenerator(std::string type) : type_(std::move(type)) {}

-  MultiSideSource<> SideOut(const std::string& tag) {
+  MultiSideSource<> SideOut(absl::string_view tag) {
    return MultiSideSource<>(&out_sides_[tag]);
  }

-  MultiSideDestination<> SideIn(const std::string& tag) {
+  MultiSideDestination<> SideIn(absl::string_view tag) {
    return MultiSideDestination<>(&in_sides_[tag]);
  }

@@ -452,19 +464,19 @@ class Graph {
  }

  // Graph ports, non-typed.
-  MultiSource<> In(const std::string& graph_input) {
+  MultiSource<> In(absl::string_view graph_input) {
    return graph_boundary_.Out(graph_input);
  }

-  MultiDestination<> Out(const std::string& graph_output) {
+  MultiDestination<> Out(absl::string_view graph_output) {
    return graph_boundary_.In(graph_output);
  }

-  MultiSideSource<> SideIn(const std::string& graph_input) {
+  MultiSideSource<> SideIn(absl::string_view graph_input) {
    return graph_boundary_.SideOut(graph_input);
  }

-  MultiSideDestination<> SideOut(const std::string& graph_output) {
+  MultiSideDestination<> SideOut(absl::string_view graph_output) {
    return graph_boundary_.SideIn(graph_output);
  }

@@ -2,6 +2,7 @@

 #include <functional>

+#include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
 #include "mediapipe/framework/api2/node.h"
 #include "mediapipe/framework/api2/packet.h"
@@ -296,6 +297,32 @@ TEST(BuilderTest, EmptyTag) {
  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
 }

+TEST(BuilderTest, StringLikeTags) {
+  const char kA[] = "A";
+  const std::string kB = "B";
+  constexpr absl::string_view kC = "C";
+
+  builder::Graph graph;
+  auto& foo = graph.AddNode("Foo");
+  graph.In(kA).SetName("a") >> foo.In(kA);
+  graph.In(kB).SetName("b") >> foo.In(kB);
+  foo.Out(kC).SetName("c") >> graph.Out(kC);
+
+  CalculatorGraphConfig expected =
+      mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+        input_stream: "A:a"
+        input_stream: "B:b"
+        output_stream: "C:c"
+        node {
+          calculator: "Foo"
+          input_stream: "A:a"
+          input_stream: "B:b"
+          output_stream: "C:c"
+        }
+      )pb");
+  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
+}
+
 TEST(BuilderTest, GraphIndexes) {
  builder::Graph graph;
  auto& foo = graph.AddNode("Foo");
@@ -326,57 +353,91 @@ TEST(BuilderTest, GraphIndexes) {

 class AnyAndSameTypeCalculator : public NodeIntf {
 public:
-  static constexpr Input<AnyType> kAnyTypeInput{"INPUT"};
-  static constexpr Output<AnyType> kAnyTypeOutput{"ANY_OUTPUT"};
-  static constexpr Output<SameType<kAnyTypeInput>> kSameTypeOutput{
+  static constexpr Input<AnyType>::Optional kAnyTypeInput{"INPUT"};
+  static constexpr Output<AnyType>::Optional kAnyTypeOutput{"ANY_OUTPUT"};
+  static constexpr Output<SameType<kAnyTypeInput>>::Optional kSameTypeOutput{
      "SAME_OUTPUT"};
+  static constexpr Output<SameType<kSameTypeOutput>> kRecursiveSameTypeOutput{
+      "RECURSIVE_SAME_OUTPUT"};

-  static constexpr Input<int> kIntInput{"INT_INPUT"};
+  static constexpr Input<int>::Optional kIntInput{"INT_INPUT"};
  // `SameType` usage for this output is only for testing purposes.
  //
  // `SameType` is designed to work with inputs of `AnyType` and, normally, you
  // would not use `Output<SameType<kIntInput>>` in a real calculator. You
  // should write `Output<int>` instead, since the type is known.
-  static constexpr Output<SameType<kIntInput>> kSameIntOutput{
+  static constexpr Output<SameType<kIntInput>>::Optional kSameIntOutput{
      "SAME_INT_OUTPUT"};
+  static constexpr Output<SameType<kSameIntOutput>> kRecursiveSameIntOutput{
+      "RECURSIVE_SAME_INT_OUTPUT"};

-  MEDIAPIPE_NODE_INTERFACE(AnyTypeCalculator, kAnyTypeInput, kAnyTypeOutput,
-                           kSameTypeOutput);
+  MEDIAPIPE_NODE_INTERFACE(AnyAndSameTypeCalculator, kAnyTypeInput,
+                           kAnyTypeOutput, kSameTypeOutput);
 };

 TEST(BuilderTest, AnyAndSameTypeHandledProperly) {
  builder::Graph graph;
-  builder::Source<internal::Generic> any_input =
-      graph[Input<AnyType>{"GRAPH_ANY_INPUT"}];
+  builder::Source<AnyType> any_input = graph[Input<AnyType>{"GRAPH_ANY_INPUT"}];
  builder::Source<int> int_input = graph[Input<int>{"GRAPH_INT_INPUT"}];

  auto& node = graph.AddNode("AnyAndSameTypeCalculator");
  any_input >> node[AnyAndSameTypeCalculator::kAnyTypeInput];
  int_input >> node[AnyAndSameTypeCalculator::kIntInput];

-  builder::Source<internal::Generic> any_type_output =
+  builder::Source<AnyType> any_type_output =
      node[AnyAndSameTypeCalculator::kAnyTypeOutput];
  any_type_output.SetName("any_type_output");

-  builder::Source<internal::Generic> same_type_output =
+  builder::Source<AnyType> same_type_output =
      node[AnyAndSameTypeCalculator::kSameTypeOutput];
  same_type_output.SetName("same_type_output");
-  builder::Source<internal::Generic> same_int_output =
+  builder::Source<AnyType> recursive_same_type_output =
+      node[AnyAndSameTypeCalculator::kRecursiveSameTypeOutput];
+  recursive_same_type_output.SetName("recursive_same_type_output");
+  builder::Source<int> same_int_output =
      node[AnyAndSameTypeCalculator::kSameIntOutput];
  same_int_output.SetName("same_int_output");
+  builder::Source<int> recursive_same_int_type_output =
+      node[AnyAndSameTypeCalculator::kRecursiveSameIntOutput];
+  recursive_same_int_type_output.SetName("recursive_same_int_type_output");
+
+  CalculatorGraphConfig expected = mediapipe::ParseTextProtoOrDie<
+      CalculatorGraphConfig>(R"pb(
+    node {
+      calculator: "AnyAndSameTypeCalculator"
+      input_stream: "INPUT:__stream_0"
+      input_stream: "INT_INPUT:__stream_1"
+      output_stream: "ANY_OUTPUT:any_type_output"
+      output_stream: "RECURSIVE_SAME_INT_OUTPUT:recursive_same_int_type_output"
+      output_stream: "RECURSIVE_SAME_OUTPUT:recursive_same_type_output"
+      output_stream: "SAME_INT_OUTPUT:same_int_output"
+      output_stream: "SAME_OUTPUT:same_type_output"
+    }
+    input_stream: "GRAPH_ANY_INPUT:__stream_0"
+    input_stream: "GRAPH_INT_INPUT:__stream_1"
+  )pb");
+  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
+}
+
+TEST(BuilderTest, AnyTypeCanBeCast) {
+  builder::Graph graph;
+  builder::Source<std::string> any_input =
+      graph.In("GRAPH_ANY_INPUT").Cast<std::string>();
+
+  auto& node = graph.AddNode("AnyAndSameTypeCalculator");
+  any_input >> node[AnyAndSameTypeCalculator::kAnyTypeInput];
+  builder::Source<double> any_type_output =
+      node[AnyAndSameTypeCalculator::kAnyTypeOutput].Cast<double>();
+  any_type_output.SetName("any_type_output");

  CalculatorGraphConfig expected =
      mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
        node {
          calculator: "AnyAndSameTypeCalculator"
          input_stream: "INPUT:__stream_0"
-          input_stream: "INT_INPUT:__stream_1"
          output_stream: "ANY_OUTPUT:any_type_output"
-          output_stream: "SAME_INT_OUTPUT:same_int_output"
-          output_stream: "SAME_OUTPUT:same_type_output"
        }
        input_stream: "GRAPH_ANY_INPUT:__stream_0"
-        input_stream: "GRAPH_INT_INPUT:__stream_1"
      )pb");
  EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
 }
@@ -27,9 +27,7 @@ using HolderBase = mediapipe::packet_internal::HolderBase;
 template <typename T>
 class Packet;

-struct DynamicType {};
-
-struct AnyType : public DynamicType {
+struct AnyType {
  AnyType() = delete;
 };

@@ -73,14 +73,12 @@ class SideOutputBase : public PortBase {
 };

 struct NoneType {
- private:
  NoneType() = delete;
 };

-template <auto& P>
-class SameType : public DynamicType {
- public:
-  static constexpr const decltype(P)& kPort = P;
+template <auto& kP>
+struct SameType {
+  static constexpr const decltype(kP)& kPort = kP;
 };

 class PacketTypeAccess;
@@ -137,21 +135,28 @@ struct IsOneOf : std::false_type {};
 template <class... T>
 struct IsOneOf<OneOf<T...>> : std::true_type {};

-template <typename T, typename std::enable_if<
-                          !std::is_base_of<DynamicType, T>{} && !IsOneOf<T>{},
-                          int>::type = 0>
+template <class T>
+struct IsSameType : std::false_type {};
+
+template <class P, P& kP>
+struct IsSameType<SameType<kP>> : std::true_type {};
+
+template <typename T,
+          typename std::enable_if<!std::is_same<T, AnyType>{} &&
+                                      !IsOneOf<T>{} && !IsSameType<T>{},
+                                  int>::type = 0>
 inline void SetType(CalculatorContract* cc, PacketType& pt) {
  pt.Set<T>();
 }

-template <typename T, typename std::enable_if<std::is_base_of<DynamicType, T>{},
-                                              int>::type = 0>
+template <typename T, typename std::enable_if<IsSameType<T>{}, int>::type = 0>
 inline void SetType(CalculatorContract* cc, PacketType& pt) {
  pt.SetSameAs(&internal::GetCollection(cc, T::kPort).Tag(T::kPort.Tag()));
 }

-template <>
-inline void SetType<AnyType>(CalculatorContract* cc, PacketType& pt) {
+template <typename T,
+          typename std::enable_if<std::is_same<T, AnyType>{}, int>::type = 0>
+inline void SetType(CalculatorContract* cc, PacketType& pt) {
  pt.SetAny();
 }

@@ -289,15 +294,15 @@ struct SideBase<InputBase> {
 };

 // TODO: maybe return a PacketBase instead of a Packet<internal::Generic>?
-template <typename T, class = void>
+template <typename T, typename = void>
 struct ActualPayloadType {
  using type = T;
 };

 template <typename T>
-struct ActualPayloadType<
-    T, std::enable_if_t<std::is_base_of<DynamicType, T>{}, void>> {
-  using type = internal::Generic;
+struct ActualPayloadType<T, std::enable_if_t<IsSameType<T>{}, void>> {
+  using type = typename ActualPayloadType<
+      typename std::decay_t<decltype(T::kPort)>::value_t>::type;
 };

 }  // namespace internal
@@ -1,4 +1,4 @@
-// Copyright 2019 The MediaPipe Authors.
+// 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.
@@ -85,14 +85,8 @@ cv::Mat MatView(const ImageFrame* image) {
  const size_t steps[] = {static_cast<size_t>(image->WidthStep()),
                          static_cast<size_t>(image->ByteDepth())};
  // Use ImageFrame to initialize in-place. ImageFrame still owns memory.
-  if (steps[0] == sizes[1] * image->NumberOfChannels() * image->ByteDepth()) {
-    // Contiguous memory optimization. See b/78570764
-    return cv::Mat(dims, sizes, type, const_cast<uint8*>(image->PixelData()));
-  } else {
-    // Custom width step.
-    return cv::Mat(dims, sizes, type, const_cast<uint8*>(image->PixelData()),
-                   steps);
-  }
+  return cv::Mat(dims, sizes, type, const_cast<uint8_t*>(image->PixelData()),
+                 steps);
 }

 }  // namespace formats
@@ -1,4 +1,4 @@
-// Copyright 2019 The MediaPipe Authors.
+// 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.
@@ -1,4 +1,4 @@
-// Copyright 2019 The MediaPipe Authors.
+// 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.
@@ -21,7 +21,6 @@
 #include "mediapipe/framework/port/logging.h"

 namespace mediapipe {
-
 namespace {

 // Set image_frame to a constant per-channel pix_value.
@@ -50,8 +49,8 @@ TEST(ImageFrameOpencvTest, ConvertToMat) {
  ImageFrame frame2(ImageFormat::GRAY8, i_width, i_height);

  // Check adding constant images.
-  const uint8 frame1_val = 12;
-  const uint8 frame2_val = 34;
+  const uint8_t frame1_val = 12;
+  const uint8_t frame2_val = 34;
  SetToColor<uint8>(&frame1_val, &frame1);
  SetToColor<uint8>(&frame2_val, &frame2);
  // Get Mat wrapper around ImageFrame memory (zero copy).
@@ -77,6 +76,37 @@ TEST(ImageFrameOpencvTest, ConvertToMat) {
  EXPECT_EQ(max_loc.y, i_height - 6);
 }

+TEST(ImageFrameOpencvTest, ConvertToIpl) {
+  const int i_width = 123, i_height = 45;
+  ImageFrame frame1(ImageFormat::GRAY8, i_width, i_height);
+  ImageFrame frame2(ImageFormat::GRAY8, i_width, i_height);
+
+  // Check adding constant images.
+  const uint8_t frame1_val = 12;
+  const uint8_t frame2_val = 34;
+  SetToColor<uint8>(&frame1_val, &frame1);
+  SetToColor<uint8>(&frame2_val, &frame2);
+  const cv::Mat frame1_mat = formats::MatView(&frame1);
+  const cv::Mat frame2_mat = formats::MatView(&frame2);
+  const cv::Mat frame_sum = frame1_mat + frame2_mat;
+  const auto frame_avg = static_cast<int>(cv::mean(frame_sum).val[0]);
+  EXPECT_EQ(frame_avg, frame1_val + frame2_val);
+
+  // Check setting min/max pixels.
+  uint8* frame1_ptr = frame1.MutablePixelData();
+  frame1_ptr[(i_width - 5) + (i_height - 5) * frame1.WidthStep()] = 1;
+  frame1_ptr[(i_width - 6) + (i_height - 6) * frame1.WidthStep()] = 100;
+  double min, max;
+  cv::Point min_loc, max_loc;
+  cv::minMaxLoc(frame1_mat, &min, &max, &min_loc, &max_loc);
+  EXPECT_EQ(min, 1);
+  EXPECT_EQ(min_loc.x, i_width - 5);
+  EXPECT_EQ(min_loc.y, i_height - 5);
+  EXPECT_EQ(max, 100);
+  EXPECT_EQ(max_loc.x, i_width - 6);
+  EXPECT_EQ(max_loc.y, i_height - 6);
+}
+
 TEST(ImageFrameOpencvTest, ImageFormats) {
  const int i_width = 123, i_height = 45;
  ImageFrame frame_g8(ImageFormat::GRAY8, i_width, i_height);
@@ -1,4 +1,4 @@
-// Copyright 2019 The MediaPipe Authors.
+// 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.
@@ -1,4 +1,4 @@
-// Copyright 2019-2020 The MediaPipe Authors.
+// 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.
@@ -37,26 +37,21 @@ namespace mediapipe {
 bool IsPowerOfTwo(int v) { return (v & (v - 1)) == 0; }

 int BhwcBatchFromShape(const Tensor::Shape& shape) {
-  LOG_IF(FATAL, shape.dims.empty())
-      << "Tensor::Shape must be non-empty to retrieve a named dimension";
+  if (shape.dims.empty()) {
+    return 1;
+  }
  return shape.dims[0];
 }

 int BhwcHeightFromShape(const Tensor::Shape& shape) {
-  LOG_IF(FATAL, shape.dims.empty())
-      << "Tensor::Shape must be non-empty to retrieve a named dimension";
  return shape.dims.size() < 4 ? 1 : shape.dims[shape.dims.size() - 3];
 }

 int BhwcWidthFromShape(const Tensor::Shape& shape) {
-  LOG_IF(FATAL, shape.dims.empty())
-      << "Tensor::Shape must be non-empty to retrieve a named dimension";
  return shape.dims.size() < 3 ? 1 : shape.dims[shape.dims.size() - 2];
 }

 int BhwcDepthFromShape(const Tensor::Shape& shape) {
-  LOG_IF(FATAL, shape.dims.empty())
-      << "Tensor::Shape must be non-empty to retrieve a named dimension";
  return shape.dims.size() < 2 ? 1 : shape.dims[shape.dims.size() - 1];
 }

@@ -424,14 +419,36 @@ Tensor::Tensor(ElementType element_type, const Shape& shape,

 #if MEDIAPIPE_METAL_ENABLED
 void Tensor::Invalidate() {
-  absl::MutexLock lock(&view_mutex_);
-  // If memory is allocated and not owned by the metal buffer.
-  // TODO: Re-design cpu buffer memory management.
-  if (cpu_buffer_ && !metal_buffer_) {
-    DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
+#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
+  GLuint cleanup_gl_tex = GL_INVALID_INDEX;
+  GLuint cleanup_gl_fb = GL_INVALID_INDEX;
+#endif  // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
+  {
+    absl::MutexLock lock(&view_mutex_);
+    // If memory is allocated and not owned by the metal buffer.
+    // TODO: Re-design cpu buffer memory management.
+    if (cpu_buffer_ && !metal_buffer_) {
+      DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
+    }
+    metal_buffer_ = nil;
+    cpu_buffer_ = nullptr;
+#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
+    // Don't need to wait for the resource to be deleted bacause if will be
+    // released on last reference deletion inside the OpenGL driver.
+    std::swap(cleanup_gl_tex, opengl_texture2d_);
+    std::swap(cleanup_gl_fb, frame_buffer_);
+#endif  // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
  }
-  metal_buffer_ = nil;
-  cpu_buffer_ = nullptr;
+#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
+  // Do not hold the view mutex while invoking GlContext::RunWithoutWaiting,
+  // since that method may acquire the context's own lock.
+  if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX) {
+    gl_context_->RunWithoutWaiting([cleanup_gl_tex, cleanup_gl_fb]() {
+      glDeleteTextures(1, &cleanup_gl_tex);
+      glDeleteFramebuffers(1, &cleanup_gl_fb);
+    });
+  }
+#endif  // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
 }

 #else
@@ -18,6 +18,7 @@
 #include <algorithm>
 #include <functional>
 #include <initializer_list>
+#include <numeric>
 #include <tuple>
 #include <type_traits>
 #include <utility>
@@ -95,9 +96,8 @@ class Tensor {
    Shape(std::initializer_list<int> dimensions) : dims(dimensions) {}
    Shape(const std::vector<int>& dimensions) : dims(dimensions) {}
    int num_elements() const {
-      int res = dims.empty() ? 0 : 1;
-      std::for_each(dims.begin(), dims.end(), [&res](int i) { res *= i; });
-      return res;
+      return std::accumulate(dims.begin(), dims.end(), 1,
+                             std::multiplies<int>());
    }
    std::vector<int> dims;
  };
@@ -15,7 +15,7 @@ def mediapipe_cc_test(
        platforms = ["linux", "android", "ios", "wasm"],
        exclude_platforms = None,
        # ios_unit_test arguments
-        ios_minimum_os_version = "9.0",
+        ios_minimum_os_version = "11.0",
        # android_cc_test arguments
        open_gl_driver = None,
        emulator_mini_boot = True,
@@ -108,6 +108,7 @@ cc_library(
        ":sharded_map",
        "//mediapipe/framework:calculator_cc_proto",
        "//mediapipe/framework:calculator_profile_cc_proto",
+        "//mediapipe/framework/port:file_helpers",
        "//mediapipe/framework/port:integral_types",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/types:optional",
@@ -22,6 +22,7 @@
 #include "absl/time/time.h"
 #include "mediapipe/framework/port/advanced_proto_lite_inc.h"
 #include "mediapipe/framework/port/canonical_errors.h"
+#include "mediapipe/framework/port/file_helpers.h"
 #include "mediapipe/framework/port/logging.h"
 #include "mediapipe/framework/port/proto_ns.h"
 #include "mediapipe/framework/port/re2.h"
@@ -244,7 +245,16 @@ absl::Status GraphProfiler::Start(mediapipe::Executor* executor) {
      executor != nullptr) {
    // Inform the user via logging the path to the trace logs.
    ASSIGN_OR_RETURN(std::string trace_log_path, GetTraceLogPath());
-    LOG(INFO) << "trace_log_path: " << trace_log_path;
+    // Check that we can actually write to it.
+    auto status =
+        file::SetContents(absl::StrCat(trace_log_path, "trace_writing_check"),
+                          "can write trace logs to this location");
+    if (status.ok()) {
+      LOG(INFO) << "trace_log_path: " << trace_log_path;
+    } else {
+      LOG(ERROR) << "cannot write to trace_log_path: " << trace_log_path << ": "
+                 << status;
+    }

    is_running_ = true;
    executor->Schedule([this] {
@@ -5,7 +5,7 @@

 buffers: {
  size_kb: 150000
-  fill_policy: DISCARD
+  fill_policy: RING_BUFFER
 }

 data_sources: {
@@ -14,19 +14,21 @@ data_sources: {
  }
 }
 data_sources: {
-    config {
-        name: "linux.ftrace"
-        ftrace_config {
-            # Scheduling information & process tracking. Useful for:
-            # - what is happening on each CPU at each moment
-            ftrace_events: "power/cpu_frequency"
-            ftrace_events: "power/cpu_idle"
-            ftrace_events: "sched/sched_switch"
-            compact_sched {
-              enabled: true
-            }
-       }
+  config {
+    name: "linux.ftrace"
+    ftrace_config {
+      # Scheduling information & process tracking. Useful for:
+      # - what is happening on each CPU at each moment
+      ftrace_events: "power/cpu_frequency"
+      ftrace_events: "power/cpu_idle"
+      ftrace_events: "sched/sched_switch"
+      compact_sched {
+        enabled: true
+      }
    }
+  }
 }
 write_into_file: true
 file_write_period_ms: 500
+# b/243571696 Added to remove Perfetto timeouts when running benchmarks remotely.
+duration_ms: 60000
@@ -821,6 +821,19 @@ cc_library(
    alwayslink = 1,
 )

+mediapipe_cc_test(
+    name = "switch_demux_calculator_test",
+    srcs = ["switch_demux_calculator_test.cc"],
+    deps = [
+        ":container_util",
+        ":switch_demux_calculator",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/port:logging",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
 cc_library(
    name = "switch_mux_calculator",
    srcs = ["switch_mux_calculator.cc"],
@@ -129,12 +129,12 @@ absl::Status SwitchDemuxCalculator::Open(CalculatorContext* cc) {
  // Relay side packets to all channels.
  // Note: This is necessary because Calculator::Open only proceeds when every
  // anticipated side-packet arrives.
-  int channel_count = tool::ChannelCount(cc->OutputSidePackets().TagMap());
+  int side_channel_count = tool::ChannelCount(cc->OutputSidePackets().TagMap());
  for (const std::string& tag : ChannelTags(cc->OutputSidePackets().TagMap())) {
    int num_entries = cc->InputSidePackets().NumEntries(tag);
    for (int index = 0; index < num_entries; ++index) {
      Packet input = cc->InputSidePackets().Get(tag, index);
-      for (int channel = 0; channel < channel_count; ++channel) {
+      for (int channel = 0; channel < side_channel_count; ++channel) {
        std::string output_tag = tool::ChannelTag(tag, channel);
        auto output_id = cc->OutputSidePackets().GetId(output_tag, index);
        if (output_id.IsValid()) {
@@ -143,6 +143,23 @@ absl::Status SwitchDemuxCalculator::Open(CalculatorContext* cc) {
      }
    }
  }
+
+  // Relay headers to all channels.
+  int output_channel_count = tool::ChannelCount(cc->Outputs().TagMap());
+  for (const std::string& tag : ChannelTags(cc->Outputs().TagMap())) {
+    int num_entries = cc->Inputs().NumEntries(tag);
+    for (int index = 0; index < num_entries; ++index) {
+      auto& input = cc->Inputs().Get(tag, index);
+      if (input.Header().IsEmpty()) continue;
+      for (int channel = 0; channel < output_channel_count; ++channel) {
+        std::string output_tag = tool::ChannelTag(tag, channel);
+        auto output_id = cc->Outputs().GetId(output_tag, index);
+        if (output_id.IsValid()) {
+          cc->Outputs().Get(output_tag, index).SetHeader(input.Header());
+        }
+      }
+    }
+  }
  return absl::OkStatus();
 }

@@ -0,0 +1,135 @@
+// 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 <string>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/port/gmock.h"
+#include "mediapipe/framework/port/gtest.h"
+#include "mediapipe/framework/port/logging.h"
+#include "mediapipe/framework/tool/container_util.h"
+
+namespace mediapipe {
+namespace {
+
+// Returns a CalculatorGraph to run a single calculator.
+CalculatorGraph BuildCalculatorGraph(CalculatorGraphConfig::Node node_config) {
+  CalculatorGraphConfig config;
+  *config.add_node() = node_config;
+  *config.mutable_input_stream() = node_config.input_stream();
+  *config.mutable_output_stream() = node_config.output_stream();
+  *config.mutable_input_side_packet() = node_config.input_side_packet();
+  *config.mutable_output_side_packet() = node_config.output_side_packet();
+  return CalculatorGraph(config);
+}
+
+// Creates a string packet.
+Packet pack(std::string data, int timestamp) {
+  return MakePacket<std::string>(data).At(Timestamp(timestamp));
+}
+
+// Creates an int packet.
+Packet pack(int data, int timestamp) {
+  return MakePacket<int>(data).At(Timestamp(timestamp));
+}
+
+// Tests showing packet channel synchronization through SwitchDemuxCalculator.
+class SwitchDemuxCalculatorTest : public ::testing::Test {
+ protected:
+  SwitchDemuxCalculatorTest() {}
+  ~SwitchDemuxCalculatorTest() override {}
+  void SetUp() override {}
+  void TearDown() override {}
+
+  // Defines a SwitchDemuxCalculator CalculatorGraphConfig::Node.
+  CalculatorGraphConfig::Node BuildNodeConfig() {
+    CalculatorGraphConfig::Node result;
+    *result.mutable_calculator() = "SwitchDemuxCalculator";
+    *result.add_input_stream() = "SELECT:select";
+    for (int c = 0; c < 2; ++c) {
+      *result.add_output_stream() =
+          absl::StrCat(tool::ChannelTag("FRAME", c), ":frame_", c);
+      *result.add_output_stream() =
+          absl::StrCat(tool::ChannelTag("MASK", c), ":mask_", c);
+    }
+    *result.add_input_stream() = "FRAME:frame";
+    *result.add_input_stream() = "MASK:mask";
+    return result;
+  }
+};
+
+// Shows the SwitchMuxCalculator is available.
+TEST_F(SwitchDemuxCalculatorTest, IsRegistered) {
+  EXPECT_TRUE(CalculatorBaseRegistry::IsRegistered("SwitchDemuxCalculator"));
+}
+
+TEST_F(SwitchDemuxCalculatorTest, BasicDataFlow) {
+  CalculatorGraphConfig::Node node_config = BuildNodeConfig();
+  CalculatorGraph graph = BuildCalculatorGraph(node_config);
+  std::vector<Packet> output_frames0;
+  EXPECT_TRUE(graph
+                  .ObserveOutputStream("frame_0",
+                                       [&](const Packet& p) {
+                                         output_frames0.push_back(p);
+                                         return absl::OkStatus();
+                                       })
+                  .ok());
+  std::vector<Packet> output_frames1;
+  EXPECT_TRUE(graph
+                  .ObserveOutputStream("frame_1",
+                                       [&](const Packet& p) {
+                                         output_frames1.push_back(p);
+                                         return absl::OkStatus();
+                                       })
+                  .ok());
+  EXPECT_TRUE(
+      graph.StartRun({}, {{"frame", MakePacket<std::string>("frame_header")}})
+          .ok());
+
+  // Finalize input for the "mask" input stream.
+  EXPECT_TRUE(graph.CloseInputStream("mask").ok());
+
+  // Channel 0 is selected just before corresponding packets arrive.
+  EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(0, 1)).ok());
+  EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(0, 10)).ok());
+  EXPECT_TRUE(graph.AddPacketToInputStream("frame", pack("p0_t10", 10)).ok());
+  EXPECT_TRUE(graph.WaitUntilIdle().ok());
+  EXPECT_EQ(output_frames0.size(), 1);
+  EXPECT_EQ(output_frames1.size(), 0);
+  EXPECT_EQ(output_frames0[0].Get<std::string>(), "p0_t10");
+
+  // Channel 1 is selected just before corresponding packets arrive.
+  EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(1, 11)).ok());
+  EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(1, 20)).ok());
+  EXPECT_TRUE(graph.AddPacketToInputStream("frame", pack("p1_t20", 20)).ok());
+  EXPECT_TRUE(graph.WaitUntilIdle().ok());
+  EXPECT_EQ(output_frames0.size(), 1);
+  EXPECT_EQ(output_frames1.size(), 1);
+  EXPECT_EQ(output_frames1[0].Get<std::string>(), "p1_t20");
+
+  EXPECT_EQ(
+      graph.FindOutputStreamManager("frame_0")->Header().Get<std::string>(),
+      "frame_header");
+  EXPECT_EQ(
+      graph.FindOutputStreamManager("frame_1")->Header().Get<std::string>(),
+      "frame_header");
+
+  EXPECT_TRUE(graph.CloseAllPacketSources().ok());
+  EXPECT_TRUE(graph.WaitUntilDone().ok());
+}
+
+}  // namespace
+}  // namespace mediapipe