Project import generated by Copybara.

GitOrigin-RevId: 73d686c40057684f8bfaca285368bf1813f9fc26

mediapipe/framework/BUILD

@@ -150,6 +150,13 @@ mediapipe_proto_library(
     deps = ["//mediapipe/framework:mediapipe_options_proto"],
 )
 
+config_setting(
+    name = "android_no_jni",
+    define_values = {"MEDIAPIPE_NO_JNI": "1"},
+    values = {"crosstool_top": "//external:android/crosstool"},
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "calculator_base",
     srcs = ["calculator_base.cc"],
@@ -712,6 +719,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/synchronization",
     ],
 )
 
@@ -916,15 +924,19 @@ cc_library(
         ":packet",
         ":packet_set",
         ":type_map",
+        "//mediapipe/framework/deps:no_destructor",
         "//mediapipe/framework/port:logging",
         "//mediapipe/framework/port:map_util",
-        "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:source_location",
         "//mediapipe/framework/port:status",
         "//mediapipe/framework/tool:status_util",
+        "//mediapipe/framework/tool:type_util",
         "//mediapipe/framework/tool:validate_name",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/types:variant",
     ],
 )
 

mediapipe/framework/api2/BUILD

@@ -134,6 +134,7 @@ cc_test(
     deps = [
         ":packet",
         "//mediapipe/framework/port:gtest_main",
+        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
     ],
 )

mediapipe/framework/api2/builder.h

@@ -313,8 +313,8 @@ template <class Calc>
 class Node : public NodeBase {
  public:
   Node() : NodeBase(Calc::kCalculatorName) {}
-  // Overrides the built-in calculator type std::string with the provided
-  // argument. Can be used to create nodes from pure interfaces.
+  // Overrides the built-in calculator type string with the provided argument.
+  // Can be used to create nodes from pure interfaces.
   // TODO: only use this for pure interfaces
   Node(const std::string& type_override) : NodeBase(type_override) {}
 
@@ -377,6 +377,29 @@ class PacketGenerator {
     return *options_.MutableExtension(T::ext);
   }
 
+  template <typename B, typename T, bool kIsOptional, bool kIsMultiple>
+  auto operator[](const PortCommon<B, T, kIsOptional, kIsMultiple>& port) {
+    using PayloadT =
+        typename PortCommon<B, T, kIsOptional, kIsMultiple>::PayloadT;
+    if constexpr (std::is_same_v<B, SideOutputBase>) {
+      auto* base = &out_sides_[port.Tag()];
+      if constexpr (kIsMultiple) {
+        return MultiSideSource<PayloadT>(base);
+      } else {
+        return SideSource<PayloadT>(base);
+      }
+    } else if constexpr (std::is_same_v<B, SideInputBase>) {
+      auto* base = &in_sides_[port.Tag()];
+      if constexpr (kIsMultiple) {
+        return MultiSideDestination<PayloadT>(base);
+      } else {
+        return SideDestination<PayloadT>(base);
+      }
+    } else {
+      static_assert(dependent_false<B>::value, "Type not supported.");
+    }
+  }
+
  private:
   std::string type_;
   TagIndexMap<DestinationBase> in_sides_;
@@ -402,7 +425,7 @@ class Graph {
   }
 
   // Creates a node of a specific type. Should be used for pure interfaces,
-  // which do not have a built-in type std::string.
+  // which do not have a built-in type string.
   template <class Calc>
   Node<Calc>& AddNode(const std::string& type) {
     auto node = std::make_unique<Node<Calc>>(type);

mediapipe/framework/api2/const_str.h

@@ -6,8 +6,8 @@
 namespace mediapipe {
 namespace api2 {
 
-// This class stores a constant std::string that can be inspected at compile
-// time in constexpr code.
+// This class stores a constant string that can be inspected at compile time
+// in constexpr code.
 class const_str {
  public:
   constexpr const_str(std::size_t size, const char* data)

mediapipe/framework/api2/packet.h

@@ -215,6 +215,7 @@ class Packet : public Packet<internal::Generic> {
     return typed_payload->data();
   }
   const T& operator*() const { return Get(); }
+  const T* operator->() const { return &Get(); }
 
   template <typename U>
   T GetOr(U&& v) const {

mediapipe/framework/api2/packet_test.cc

@@ -1,5 +1,6 @@
 #include "mediapipe/framework/api2/packet.h"
 
+#include "absl/memory/memory.h"
 #include "absl/strings/str_cat.h"
 #include "mediapipe/framework/port/gmock.h"
 #include "mediapipe/framework/port/gtest.h"
@@ -18,6 +19,17 @@ class LiveCheck {
   bool& alive_;
 };
 
+class Base {
+ public:
+  virtual ~Base() = default;
+  virtual absl::string_view name() const { return "Base"; }
+};
+
+class Derived : public Base {
+ public:
+  absl::string_view name() const override { return "Derived"; }
+};
+
 TEST(PacketTest, PacketBaseDefault) {
   PacketBase p;
   EXPECT_TRUE(p.IsEmpty());
@@ -242,6 +254,16 @@ TEST(PacketTest, OneOfConsume) {
   EXPECT_TRUE(p.IsEmpty());
 }
 
+TEST(PacketTest, Polymorphism) {
+  Packet<Base> base = PacketAdopting<Base>(absl::make_unique<Derived>());
+  EXPECT_EQ(base->name(), "Derived");
+  // Since packet contents are implicitly immutable, if you need mutability the
+  // current recommendation is still to wrap the contents in a unique_ptr.
+  Packet<std::unique_ptr<Base>> mutable_base =
+      MakePacket<std::unique_ptr<Base>>(absl::make_unique<Derived>());
+  EXPECT_EQ((**mutable_base).name(), "Derived");
+}
+
 }  // namespace
 }  // namespace api2
 }  // namespace mediapipe

mediapipe/framework/api2/port.h

@@ -172,9 +172,14 @@ inline void SetType<NoneType>(CalculatorContract* cc, PacketType& pt) {
   pt.SetNone();
 }
 
+template <typename... T>
+inline void SetTypeOneOf(OneOf<T...>, CalculatorContract* cc, PacketType& pt) {
+  pt.SetOneOf<T...>();
+}
+
 template <typename T, typename std::enable_if<IsOneOf<T>{}, int>::type = 0>
 inline void SetType(CalculatorContract* cc, PacketType& pt) {
-  pt.SetAny();
+  SetTypeOneOf(T{}, cc, pt);
 }
 
 template <typename ValueT>
@@ -294,14 +299,26 @@ struct SideBase<InputBase> {
   using type = SideInputBase;
 };
 
+// TODO: maybe return a PacketBase instead of a Packet<internal::Generic>?
+template <typename T, class = 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;
+};
+
 }  // namespace internal
 
-// TODO: maybe return a PacketBase instead of a Packet<internal::Generic>?
-template <typename T, typename std::enable_if<
-                          !std::is_base_of<DynamicType, T>{}, int>::type = 0>
-auto ActualValueT(T) -> T;
+// Maps special port value types, such as AnyType, to internal::Generic.
+template <typename T>
+using ActualPayloadT = typename internal::ActualPayloadType<T>::type;
 
-auto ActualValueT(DynamicType) -> internal::Generic;
+static_assert(std::is_same_v<ActualPayloadT<int>, int>, "");
+static_assert(std::is_same_v<ActualPayloadT<AnyType>, internal::Generic>, "");
 
 template <typename Base, typename ValueT, bool IsOptional = false,
           bool IsMultiple = false>
@@ -325,7 +342,7 @@ class PortCommon : public Base {
   explicit constexpr PortCommon(const char (&tag)[N])
       : Base(N, tag, &get_type_hash<ValueT>, IsOptionalV, IsMultipleV) {}
 
-  using PayloadT = decltype(ActualValueT(std::declval<ValueT>()));
+  using PayloadT = ActualPayloadT<ValueT>;
 
   auto operator()(CalculatorContext* cc) const {
     return internal::AccessPort<PayloadT>(
@@ -385,7 +402,7 @@ class SideFallbackT : public Base {
   static constexpr bool kOptional = IsOptionalV;
   static constexpr bool kMultiple = IsMultipleV;
   using Optional = SideFallbackT<Base, ValueT, true, IsMultipleV>;
-  using PayloadT = decltype(ActualValueT(std::declval<ValueT>()));
+  using PayloadT = ActualPayloadT<ValueT>;
 
   const char* Tag() const { return stream_port.Tag(); }
 
@@ -499,6 +516,10 @@ class OutputShardAccess : public OutputShardAccessBase {
     Send(std::move(payload), context_.InputTimestamp());
   }
 
+  void SetHeader(const PacketBase& header) {
+    if (output_) output_->SetHeader(ToOldPacket(header));
+  }
+
  private:
   OutputShardAccess(const CalculatorContext& cc, OutputStreamShard* output)
       : OutputShardAccessBase(cc, output) {}

mediapipe/framework/api2/port_test.cc

@@ -21,6 +21,25 @@ TEST(PortTest, Tag) {
   EXPECT_EQ(std::string(port.Tag()), "FOO");
 }
 
+struct DeletedCopyType {
+  DeletedCopyType(const DeletedCopyType&) = delete;
+  DeletedCopyType& operator=(const DeletedCopyType&) = delete;
+};
+
+TEST(PortTest, DeletedCopyConstructorInput) {
+  static constexpr Input<DeletedCopyType> kInputPort{"INPUT"};
+  EXPECT_EQ(std::string(kInputPort.Tag()), "INPUT");
+
+  static constexpr Output<DeletedCopyType> kOutputPort{"OUTPUT"};
+  EXPECT_EQ(std::string(kOutputPort.Tag()), "OUTPUT");
+
+  static constexpr SideInput<DeletedCopyType> kSideInputPort{"SIDE_INPUT"};
+  EXPECT_EQ(std::string(kSideInputPort.Tag()), "SIDE_INPUT");
+
+  static constexpr SideOutput<DeletedCopyType> kSideOutputPort{"SIDE_OUTPUT"};
+  EXPECT_EQ(std::string(kSideOutputPort.Tag()), "SIDE_OUTPUT");
+}
+
 }  // namespace
 }  // namespace api2
 }  // namespace mediapipe

mediapipe/framework/api2/tag_test.cc

@@ -26,8 +26,8 @@ TEST(TagTest, String) {
   EXPECT_EQ(kBAR.str(), "BAR");
 }
 
-// Separate invocations of MPP_TAG with the same std::string produce objects of
-// the same type.
+// Separate invocations of MPP_TAG with the same string produce objects of the
+// same type.
 TEST(TagTest, SameType) { EXPECT_TRUE(same_type(kFOO, kFOO2)); }
 
 // Different tags have different types.

mediapipe/framework/calculator_contract.h

@@ -95,8 +95,8 @@ class CalculatorContract {
     input_stream_handler_options_ = options;
   }
 
-  // Returns the name of this Nodes's InputStreamHandler, or empty std::string
-  // if none is set.
+  // Returns the name of this Nodes's InputStreamHandler, or empty string if
+  // none is set.
   std::string GetInputStreamHandler() const { return input_stream_handler_; }
 
   // Returns the MediaPipeOptions of this Node's InputStreamHandler, or empty

mediapipe/framework/calculator_graph.h

@@ -54,15 +54,13 @@
 #include "mediapipe/framework/scheduler.h"
 #include "mediapipe/framework/thread_pool_executor.pb.h"
 
-#if !MEDIAPIPE_DISABLE_GPU
 namespace mediapipe {
+
+#if !MEDIAPIPE_DISABLE_GPU
 class GpuResources;
 struct GpuSharedData;
-}  // namespace mediapipe
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
-namespace mediapipe {
-
 typedef absl::StatusOr<OutputStreamPoller> StatusOrPoller;
 
 // The class representing a DAG of calculator nodes.
@@ -366,10 +364,9 @@ class CalculatorGraph {
 #if !MEDIAPIPE_DISABLE_GPU
   // Returns a pointer to the GpuResources in use, if any.
   // Only meant for internal use.
-  std::shared_ptr<::mediapipe::GpuResources> GetGpuResources() const;
+  std::shared_ptr<GpuResources> GetGpuResources() const;
 
-  absl::Status SetGpuResources(
-      std::shared_ptr<::mediapipe::GpuResources> resources);
+  absl::Status SetGpuResources(std::shared_ptr<GpuResources> resources);
 
   // Helper for PrepareForRun. If it returns a non-empty map, those packets
   // must be added to the existing side packets, replacing existing values
@@ -532,7 +529,7 @@ class CalculatorGraph {
 #if !MEDIAPIPE_DISABLE_GPU
   // Owns the legacy GpuSharedData if we need to create one for backwards
   // compatibility.
-  std::unique_ptr<::mediapipe::GpuSharedData> legacy_gpu_shared_;
+  std::unique_ptr<GpuSharedData> legacy_gpu_shared_;
 #endif  // !MEDIAPIPE_DISABLE_GPU
 
   // True if the graph was initialized.
@@ -598,7 +595,7 @@ class CalculatorGraph {
   std::unique_ptr<CounterFactory> counter_factory_;
 
   // Executors for the scheduler, keyed by the executor's name. The default
-  // executor's name is the empty std::string.
+  // executor's name is the empty string.
   std::map<std::string, std::shared_ptr<Executor>> executors_;
 
   // The processed input side packet map for this run.

mediapipe/framework/calculator_graph_test.cc

@@ -768,7 +768,7 @@ typedef TypedStatusHandler<uint32> Uint32StatusHandler;
 REGISTER_STATUS_HANDLER(StringStatusHandler);
 REGISTER_STATUS_HANDLER(Uint32StatusHandler);
 
-// A std::string generator that will succeed.
+// A string generator that will succeed.
 class StaticCounterStringGenerator : public PacketGenerator {
  public:
   static absl::Status FillExpectations(
@@ -1767,15 +1767,14 @@ TEST(CalculatorGraph, StatusHandlerInputVerification) {
   EXPECT_FALSE(graph->Run({{"a_uint64", a_uint64}}).ok());
 
   // Should fail verification when the type of an already created packet is
-  // wrong. Here we give the uint64 packet instead of the std::string packet to
-  // the StringStatusHandler.
+  // wrong. Here we give the uint64 packet instead of the string packet to the
+  // StringStatusHandler.
   EXPECT_FALSE(
       graph->Run({{"extra_string", a_uint64}, {"a_uint64", a_uint64}}).ok());
 
   // Should fail verification when the type of a packet generated by a base
   // packet factory is wrong. Everything is correct except we add a status
-  // handler expecting a uint32 but give it the std::string from the packet
-  // factory.
+  // handler expecting a uint32 but give it the string from the packet factory.
   auto* invalid_handler = config.add_status_handler();
   invalid_handler->set_status_handler("Uint32StatusHandler");
   invalid_handler->add_input_side_packet("created_by_factory");
@@ -1792,8 +1791,8 @@ TEST(CalculatorGraph, StatusHandlerInputVerification) {
                                  MediaPipeTypeStringOrDemangled<uint32>())));
 
   // Should fail verification when the type of a to-be-generated packet is
-  // wrong. The added handler now expects a std::string but will receive the
-  // uint32 generated by the existing generator.
+  // wrong. The added handler now expects a string but will receive the uint32
+  // generated by the existing generator.
   invalid_handler->set_status_handler("StringStatusHandler");
   invalid_handler->set_input_side_packet(0, "generated_by_generator");
   graph.reset(new CalculatorGraph());

mediapipe/framework/calculator_node.h

@@ -79,10 +79,9 @@ class CalculatorNode {
   // running first. If a node is not a source, this method is not called.
   Timestamp SourceProcessOrder(const CalculatorContext* cc) const;
 
-  // Retrieves a std::string name for the node.  If the node's name was set in
-  // the calculator graph config, it will be returned.  Otherwise, a
-  // human-readable std::string that uniquely identifies the node is returned,
-  // e.g.
+  // Retrieves a string name for the node.  If the node's name was set in the
+  // calculator graph config, it will be returned.  Otherwise, a human-readable
+  // string that uniquely identifies the node is returned, e.g.
   // "[FooBarCalculator with first output stream \"foo_bar_output\"]" for
   // non-sink nodes and "[FooBarCalculator with node ID: 42 and input streams:
   // \"foo_bar_input\"]" for sink nodes.  This name should be used in error
@@ -278,7 +277,7 @@ class CalculatorNode {
   void CloseInputStreams() ABSL_LOCKS_EXCLUDED(status_mutex_);
   void CloseOutputStreams(OutputStreamShardSet* outputs)
       ABSL_LOCKS_EXCLUDED(status_mutex_);
-  // Get a std::string describing the input streams.
+  // Get a string describing the input streams.
   std::string DebugInputStreamNames() const;
 
   // Returns true if all outputs will be identical to the previous graph run.

mediapipe/framework/calculator_runner.h

@@ -62,7 +62,7 @@ class CalculatorRunner {
   //   )");
   explicit CalculatorRunner(const CalculatorGraphConfig::Node& node_config);
 #if !defined(MEDIAPIPE_PROTO_LITE)
-  // Convenience constructor which takes a node_config std::string directly.
+  // Convenience constructor which takes a node_config string directly.
   explicit CalculatorRunner(const std::string& node_config_string);
   // Convenience constructor to initialize a calculator which uses indexes
   // (not tags) for all its fields.

mediapipe/framework/deps/file_path.h

@@ -51,7 +51,8 @@ std::string JoinPathImpl(bool honor_abs,
 //
 // Usage:
 // std::string path = file::JoinPath("/cns", dirname, filename);
-// std::string path = file::JoinPath("./", filename);
+// std::string path = file::JoinPath("./",
+//                                   filename);
 //
 // 0, 1, 2-path specializations exist to optimize common cases.
 inline std::string JoinPath() { return std::string(); }
@@ -69,7 +70,7 @@ inline std::string JoinPath(absl::string_view path1, absl::string_view path2,
 // * If there is a single leading "/" in the path, the result will be the
 //   leading "/".
 // * If there is no "/" in the path, the result is the empty prefix of the
-//   input std::string.
+//   input string.
 absl::string_view Dirname(absl::string_view path);
 
 // Return the parts of the path, split on the final "/".  If there is no
@@ -83,7 +84,7 @@ std::pair<absl::string_view, absl::string_view> SplitPath(
 // "/" in the path, the result is the same as the input.
 // Note that this function's behavior differs from the Unix basename
 // command if path ends with "/". For such paths, this function returns the
-// empty std::string.
+// empty string.
 absl::string_view Basename(absl::string_view path);
 
 // Returns the part of the basename of path after the final ".".  If

mediapipe/framework/deps/numbers.h

@@ -15,6 +15,8 @@
 #ifndef MEDIAPIPE_DEPS_NUMBERS_H_
 #define MEDIAPIPE_DEPS_NUMBERS_H_
 
+#include <string>
+
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
 #include "mediapipe/framework/port/integral_types.h"

mediapipe/framework/deps/rectangle.h

@@ -145,7 +145,7 @@ class Rectangle {
   void AddBorder(const T& border_size);
 
   // Debug printing.
-  friend std::ostream& operator<<<T>(std::ostream&, const Rectangle&);
+  friend std::ostream& operator<< <T>(std::ostream&, const Rectangle&);
 
  private:
   Point2<T> min_;

mediapipe/framework/deps/registration.h

@@ -370,7 +370,7 @@ class GlobalFactoryRegistry {
   GlobalFactoryRegistry() = delete;
 };
 
-// Two levels of macros are required to convert __LINE__ into a std::string
+// Two levels of macros are required to convert __LINE__ into a string
 // containing the line number.
 #define REGISTRY_STATIC_VAR_INNER(var_name, line) var_name##_##line##__
 #define REGISTRY_STATIC_VAR(var_name, line) \

mediapipe/framework/deps/singleton.h

@@ -25,7 +25,7 @@ class Singleton {
  public:
   // Returns the pointer to the singleton of type |T|.
   // This method is thread-safe.
-  static T *get() LOCKS_EXCLUDED(mu_) {
+  static T *get() ABSL_LOCKS_EXCLUDED(mu_) {
     absl::MutexLock lock(&mu_);
     if (instance_) {
       return instance_;
@@ -46,7 +46,7 @@ class Singleton {
   // cannot be recreated. However, the callers of this method responsible for
   // making sure that no other threads are accessing (or plan to access) the
   // singleton any longer.
-  static void Destruct() LOCKS_EXCLUDED(mu_) {
+  static void Destruct() ABSL_LOCKS_EXCLUDED(mu_) {
     absl::MutexLock lock(&mu_);
     T *tmp_ptr = instance_;
     instance_ = nullptr;
@@ -55,8 +55,8 @@ class Singleton {
   }
 
  private:
-  static T *instance_ GUARDED_BY(mu_);
-  static bool destroyed_ GUARDED_BY(mu_);
+  static T *instance_ ABSL_GUARDED_BY(mu_);
+  static bool destroyed_ ABSL_GUARDED_BY(mu_);
   static absl::Mutex mu_;
 };
 

mediapipe/framework/deps/source_location.h

@@ -47,7 +47,7 @@ class source_location {
   // MEDIAPIPE_LOC macro below.
   //
   // file_name must outlive all copies of the source_location
-  // object, so in practice it should be a std::string literal.
+  // object, so in practice it should be a string literal.
   constexpr source_location(std::uint_least32_t line, const char* file_name)
       : line_(line), file_name_(file_name) {}
 

mediapipe/framework/deps/status.cc

@@ -29,7 +29,7 @@ std::string* MediaPipeCheckOpHelperOutOfLine(const absl::Status& v,
   r += msg;
   r += " status: ";
   r += v.ToString();
-  // Leaks std::string but this is only to be used in a fatal error message
+  // Leaks string but this is only to be used in a fatal error message
   return new std::string(r);
 }
 

mediapipe/framework/deps/status_builder.cc

@@ -15,6 +15,7 @@
 #include "mediapipe/framework/deps/status_builder.h"
 
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 
 namespace mediapipe {
 
@@ -23,7 +24,9 @@ StatusBuilder::StatusBuilder(const StatusBuilder& sb) {
   file_ = sb.file_;
   line_ = sb.line_;
   no_logging_ = sb.no_logging_;
-  stream_ = absl::make_unique<std::ostringstream>(sb.stream_->str());
+  stream_ = sb.stream_
+                ? absl::make_unique<std::ostringstream>(sb.stream_->str())
+                : nullptr;
   join_style_ = sb.join_style_;
 }
 
@@ -32,43 +35,58 @@ StatusBuilder& StatusBuilder::operator=(const StatusBuilder& sb) {
   file_ = sb.file_;
   line_ = sb.line_;
   no_logging_ = sb.no_logging_;
-  stream_ = absl::make_unique<std::ostringstream>(sb.stream_->str());
+  stream_ = sb.stream_
+                ? absl::make_unique<std::ostringstream>(sb.stream_->str())
+                : nullptr;
   join_style_ = sb.join_style_;
   return *this;
 }
 
-StatusBuilder& StatusBuilder::SetAppend() {
+StatusBuilder& StatusBuilder::SetAppend() & {
   if (status_.ok()) return *this;
   join_style_ = MessageJoinStyle::kAppend;
   return *this;
 }
 
-StatusBuilder& StatusBuilder::SetPrepend() {
+StatusBuilder&& StatusBuilder::SetAppend() && { return std::move(SetAppend()); }
+
+StatusBuilder& StatusBuilder::SetPrepend() & {
   if (status_.ok()) return *this;
   join_style_ = MessageJoinStyle::kPrepend;
   return *this;
 }
 
-StatusBuilder& StatusBuilder::SetNoLogging() {
+StatusBuilder&& StatusBuilder::SetPrepend() && {
+  return std::move(SetPrepend());
+}
+
+StatusBuilder& StatusBuilder::SetNoLogging() & {
   no_logging_ = true;
   return *this;
 }
 
+StatusBuilder&& StatusBuilder::SetNoLogging() && {
+  return std::move(SetNoLogging());
+}
+
 StatusBuilder::operator Status() const& {
-  if (stream_->str().empty() || no_logging_) {
+  if (!stream_ || stream_->str().empty() || no_logging_) {
     return status_;
   }
   return StatusBuilder(*this).JoinMessageToStatus();
 }
 
 StatusBuilder::operator Status() && {
-  if (stream_->str().empty() || no_logging_) {
+  if (!stream_ || stream_->str().empty() || no_logging_) {
     return status_;
   }
   return JoinMessageToStatus();
 }
 
 absl::Status StatusBuilder::JoinMessageToStatus() {
+  if (!stream_) {
+    return absl::OkStatus();
+  }
   std::string message;
   if (join_style_ == MessageJoinStyle::kAnnotate) {
     if (!status_.ok()) {

mediapipe/framework/deps/status_builder.h

@@ -15,7 +15,13 @@
 #ifndef MEDIAPIPE_DEPS_STATUS_BUILDER_H_
 #define MEDIAPIPE_DEPS_STATUS_BUILDER_H_
 
+#include <memory>
+#include <sstream>
+#include <utility>
+
 #include "absl/base/attributes.h"
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "mediapipe/framework/deps/source_location.h"
@@ -27,6 +33,10 @@ class ABSL_MUST_USE_RESULT StatusBuilder {
  public:
   StatusBuilder(const StatusBuilder& sb);
   StatusBuilder& operator=(const StatusBuilder& sb);
+
+  StatusBuilder(StatusBuilder&&) = default;
+  StatusBuilder& operator=(StatusBuilder&&) = default;
+
   // Creates a `StatusBuilder` based on an original status.  If logging is
   // enabled, it will use `location` as the location from which the log message
   // occurs.  A typical user will call this with `MEDIAPIPE_LOC`.
@@ -35,14 +45,14 @@ class ABSL_MUST_USE_RESULT StatusBuilder {
       : status_(original_status),
         line_(location.line()),
         file_(location.file_name()),
-        stream_(new std::ostringstream) {}
+        stream_(InitStream(status_)) {}
 
   StatusBuilder(absl::Status&& original_status,
                 mediapipe::source_location location)
       : status_(std::move(original_status)),
         line_(location.line()),
         file_(location.file_name()),
-        stream_(new std::ostringstream) {}
+        stream_(InitStream(status_)) {}
 
   // Creates a `StatusBuilder` from a mediapipe status code.  If logging is
   // enabled, it will use `location` as the location from which the log message
@@ -51,29 +61,37 @@ class ABSL_MUST_USE_RESULT StatusBuilder {
       : status_(code, ""),
         line_(location.line()),
         file_(location.file_name()),
-        stream_(new std::ostringstream) {}
+        stream_(InitStream(status_)) {}
 
   StatusBuilder(const absl::Status& original_status, const char* file, int line)
       : status_(original_status),
         line_(line),
         file_(file),
-        stream_(new std::ostringstream) {}
+        stream_(InitStream(status_)) {}
 
   bool ok() const { return status_.ok(); }
 
-  StatusBuilder& SetAppend();
+  StatusBuilder& SetAppend() &;
+  StatusBuilder&& SetAppend() &&;
 
-  StatusBuilder& SetPrepend();
+  StatusBuilder& SetPrepend() &;
+  StatusBuilder&& SetPrepend() &&;
 
-  StatusBuilder& SetNoLogging();
+  StatusBuilder& SetNoLogging() &;
+  StatusBuilder&& SetNoLogging() &&;
 
   template <typename T>
-  StatusBuilder& operator<<(const T& msg) {
-    if (status_.ok()) return *this;
+  StatusBuilder& operator<<(const T& msg) & {
+    if (!stream_) return *this;
     *stream_ << msg;
     return *this;
   }
 
+  template <typename T>
+  StatusBuilder&& operator<<(const T& msg) && {
+    return std::move(*this << msg);
+  }
+
   operator Status() const&;
   operator Status() &&;
 
@@ -88,6 +106,15 @@ class ABSL_MUST_USE_RESULT StatusBuilder {
     kPrepend,
   };
 
+  // Conditionally creates an ostringstream if the status is not ok.
+  static std::unique_ptr<std::ostringstream> InitStream(
+      const absl::Status status) {
+    if (status.ok()) {
+      return nullptr;
+    }
+    return absl::make_unique<std::ostringstream>();
+  }
+
   // The status that the result will be based on.
   absl::Status status_;
   // The line to record if this file is logged.
@@ -95,7 +122,8 @@ class ABSL_MUST_USE_RESULT StatusBuilder {
   // Not-owned: The file to record if this status is logged.
   const char* file_;
   bool no_logging_ = false;
-  // The additional messages added with `<<`.
+  // The additional messages added with `<<`.  This is nullptr when status_ is
+  // ok.
   std::unique_ptr<std::ostringstream> stream_;
   // Specifies how to join the message in `status_` and `stream_`.
   MessageJoinStyle join_style_ = MessageJoinStyle::kAnnotate;

mediapipe/framework/deps/status_builder_test.cc

@@ -18,6 +18,21 @@
 
 namespace mediapipe {
 
+TEST(StatusBuilder, OkStatusLvalue) {
+  StatusBuilder builder(absl::OkStatus(), MEDIAPIPE_LOC);
+  builder << "annotated message1 "
+          << "annotated message2";
+  absl::Status status = builder;
+  ASSERT_EQ(status, absl::OkStatus());
+}
+
+TEST(StatusBuilder, OkStatusRvalue) {
+  absl::Status status = StatusBuilder(absl::OkStatus(), MEDIAPIPE_LOC)
+                        << "annotated message1 "
+                        << "annotated message2";
+  ASSERT_EQ(status, absl::OkStatus());
+}
+
 TEST(StatusBuilder, AnnotateMode) {
   absl::Status status = StatusBuilder(absl::Status(absl::StatusCode::kNotFound,
                                                    "original message"),
@@ -30,7 +45,12 @@ TEST(StatusBuilder, AnnotateMode) {
             "original message; annotated message1 annotated message2");
 }
 
-TEST(StatusBuilder, PrependMode) {
+TEST(StatusBuilder, PrependModeLvalue) {
+  StatusBuilder builder(
+      absl::Status(absl::StatusCode::kInvalidArgument, "original message"),
+      MEDIAPIPE_LOC);
+  builder.SetPrepend() << "prepended message1 "
+                       << "prepended message2 ";
   absl::Status status =
       StatusBuilder(
           absl::Status(absl::StatusCode::kInvalidArgument, "original message"),
@@ -44,7 +64,33 @@ TEST(StatusBuilder, PrependMode) {
             "prepended message1 prepended message2 original message");
 }
 
-TEST(StatusBuilder, AppendMode) {
+TEST(StatusBuilder, PrependModeRvalue) {
+  absl::Status status =
+      StatusBuilder(
+          absl::Status(absl::StatusCode::kInvalidArgument, "original message"),
+          MEDIAPIPE_LOC)
+          .SetPrepend()
+      << "prepended message1 "
+      << "prepended message2 ";
+  ASSERT_FALSE(status.ok());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(),
+            "prepended message1 prepended message2 original message");
+}
+
+TEST(StatusBuilder, AppendModeLvalue) {
+  StatusBuilder builder(
+      absl::Status(absl::StatusCode::kInternal, "original message"),
+      MEDIAPIPE_LOC);
+  builder.SetAppend() << " extra message1"
+                      << " extra message2";
+  absl::Status status = builder;
+  ASSERT_FALSE(status.ok());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInternal);
+  EXPECT_EQ(status.message(), "original message extra message1 extra message2");
+}
+
+TEST(StatusBuilder, AppendModeRvalue) {
   absl::Status status = StatusBuilder(absl::Status(absl::StatusCode::kInternal,
                                                    "original message"),
                                       MEDIAPIPE_LOC)
@@ -56,7 +102,18 @@ TEST(StatusBuilder, AppendMode) {
   EXPECT_EQ(status.message(), "original message extra message1 extra message2");
 }
 
-TEST(StatusBuilder, NoLoggingMode) {
+TEST(StatusBuilder, NoLoggingModeLvalue) {
+  StatusBuilder builder(
+      absl::Status(absl::StatusCode::kUnavailable, "original message"),
+      MEDIAPIPE_LOC);
+  builder.SetNoLogging() << " extra message";
+  absl::Status status = builder;
+  ASSERT_FALSE(status.ok());
+  EXPECT_EQ(status.code(), absl::StatusCode::kUnavailable);
+  EXPECT_EQ(status.message(), "original message");
+}
+
+TEST(StatusBuilder, NoLoggingModeRvalue) {
   absl::Status status =
       StatusBuilder(
           absl::Status(absl::StatusCode::kUnavailable, "original message"),

mediapipe/framework/deps/status_macros.h

@@ -150,21 +150,28 @@
 #define STATUS_MACROS_IMPL_GET_VARIADIC_(args) \
   STATUS_MACROS_IMPL_GET_VARIADIC_HELPER_ args
 
-#define STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_2_(lhs, rexpr) \
-  STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_3_(lhs, rexpr, std::move(_))
+#define STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_2_(lhs, rexpr)                  \
+  STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_(                                     \
+      STATUS_MACROS_IMPL_CONCAT_(_status_or_value, __LINE__), lhs, rexpr,   \
+      return mediapipe::StatusBuilder(                                      \
+          std::move(STATUS_MACROS_IMPL_CONCAT_(_status_or_value, __LINE__)) \
+              .status(),                                                    \
+          __FILE__, __LINE__))
 #define STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_3_(lhs, rexpr, error_expression) \
   STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_(                                      \
       STATUS_MACROS_IMPL_CONCAT_(_status_or_value, __LINE__), lhs, rexpr,    \
-      error_expression)
-#define STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_(statusor, lhs, rexpr,      \
-                                             error_expression)          \
-  auto statusor = (rexpr);                                              \
-  if (ABSL_PREDICT_FALSE(!statusor.ok())) {                             \
-    mediapipe::StatusBuilder _(std::move(statusor).status(), __FILE__,  \
-                               __LINE__);                               \
-    (void)_; /* error_expression is allowed to not use this variable */ \
-    return (error_expression);                                          \
-  }                                                                     \
+      mediapipe::StatusBuilder _(                                            \
+          std::move(STATUS_MACROS_IMPL_CONCAT_(_status_or_value, __LINE__))  \
+              .status(),                                                     \
+          __FILE__, __LINE__);                                               \
+      (void)_; /* error_expression is allowed to not use this variable */    \
+      return (error_expression))
+#define STATUS_MACROS_IMPL_ASSIGN_OR_RETURN_(statusor, lhs, rexpr, \
+                                             error_expression)     \
+  auto statusor = (rexpr);                                         \
+  if (ABSL_PREDICT_FALSE(!statusor.ok())) {                        \
+    error_expression;                                              \
+  }                                                                \
   lhs = std::move(statusor).value()
 
 // Internal helper for concatenating macro values.

mediapipe/framework/formats/BUILD

@@ -332,15 +332,13 @@ cc_library(
         "//mediapipe/framework:port",
         "//mediapipe/framework:type_map",
         "//mediapipe/framework/port:logging",
+        "//mediapipe/gpu:gpu_buffer",
+        "//mediapipe/gpu:gpu_buffer_format",
     ] + select({
         "//conditions:default": [
-            "//mediapipe/gpu:gpu_buffer",
-            "//mediapipe/gpu:gpu_buffer_format",
             "//mediapipe/gpu:gl_texture_buffer",
         ],
         "//mediapipe:ios": [
-            "//mediapipe/gpu:gpu_buffer",
-            "//mediapipe/gpu:gpu_buffer_format",
         ],
         "//mediapipe/gpu:disable_gpu": [],
     }) + select({
@@ -430,7 +428,10 @@ cc_test(
 
 cc_library(
     name = "tensor",
-    srcs = ["tensor.cc"],
+    srcs =
+        [
+            "tensor.cc",
+        ],
     hdrs = ["tensor.h"],
     copts = select({
         "//mediapipe:apple": [

mediapipe/framework/formats/image.cc

@@ -16,48 +16,15 @@
 
 #include "mediapipe/framework/type_map.h"
 
+#if !MEDIAPIPE_DISABLE_GPU
+#include "mediapipe/gpu/gl_texture_view.h"
+#endif  // !MEDIAPIPE_DISABLE_GPU
+
 namespace mediapipe {
 
 // TODO Refactor common code from GpuBufferToImageFrameCalculator
 bool Image::ConvertToCpu() const {
-  if (!use_gpu_) return true;  // Already on CPU.
-#if !MEDIAPIPE_DISABLE_GPU
-#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-  image_frame_ = CreateImageFrameForCVPixelBuffer(GetCVPixelBufferRef());
-#else
-  auto gl_texture = gpu_buffer_.GetGlTextureBufferSharedPtr();
-  if (!gl_texture->GetProducerContext()) return false;
-  gl_texture->GetProducerContext()->Run([this, &gl_texture]() {
-    gl_texture->WaitOnGpu();
-    const auto gpu_buf = mediapipe::GpuBuffer(GetGlTextureBufferSharedPtr());
-#ifdef __ANDROID__
-    glBindFramebuffer(GL_FRAMEBUFFER, 0);  // b/32091368
-#endif
-    GLuint fb = 0;
-    glDisable(GL_DEPTH_TEST);
-    // TODO Re-use a shared framebuffer.
-    glGenFramebuffers(1, &fb);
-    glBindFramebuffer(GL_FRAMEBUFFER, fb);
-    glViewport(0, 0, gpu_buf.width(), gpu_buf.height());
-    glActiveTexture(GL_TEXTURE0);
-    glBindTexture(gl_texture->target(), gl_texture->name());
-    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
-                           gl_texture->target(), gl_texture->name(), 0);
-    auto frame = std::make_shared<ImageFrame>(
-        mediapipe::ImageFormatForGpuBufferFormat(gpu_buf.format()),
-        gpu_buf.width(), gpu_buf.height(),
-        ImageFrame::kGlDefaultAlignmentBoundary);
-    const auto info = GlTextureInfoForGpuBufferFormat(
-        gpu_buf.format(), 0, gl_texture->GetProducerContext()->GetGlVersion());
-    glReadPixels(0, 0, gpu_buf.width(), gpu_buf.height(), info.gl_format,
-                 info.gl_type, frame->MutablePixelData());
-    glDeleteFramebuffers(1, &fb);
-    // Cleanup
-    gl_texture->DidRead(gl_texture->GetProducerContext()->CreateSyncToken());
-    image_frame_ = frame;
-  });
-#endif  //  MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-#endif  // !MEDIAPIPE_DISABLE_GPU
+  auto view = gpu_buffer_.GetReadView<ImageFrame>();
   use_gpu_ = false;
   return true;
 }
@@ -67,19 +34,7 @@ bool Image::ConvertToGpu() const {
 #if MEDIAPIPE_DISABLE_GPU
   return false;
 #else
-  if (use_gpu_) return true;  // Already on GPU.
-#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
-  auto packet = PointToForeign<ImageFrame>(image_frame_.get());
-  CFHolder<CVPixelBufferRef> buffer;
-  auto status = CreateCVPixelBufferForImageFramePacket(packet, true, &buffer);
-  CHECK_OK(status);
-  gpu_buffer_ = mediapipe::GpuBuffer(std::move(buffer));
-#else
-  // GlCalculatorHelperImpl::MakeGlTextureBuffer (CreateSourceTexture)
-  auto buffer = mediapipe::GlTextureBuffer::Create(*image_frame_);
-  glFlush();
-  gpu_buffer_ = mediapipe::GpuBuffer(std::move(buffer));
-#endif  //  MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+  auto view = gpu_buffer_.GetReadView<GlTextureView>(0);
   use_gpu_ = true;
   return true;
 #endif  // MEDIAPIPE_DISABLE_GPU

mediapipe/framework/formats/image.h

@@ -21,20 +21,18 @@
 #include "mediapipe/framework/formats/image_format.pb.h"
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/framework/port/logging.h"
-
-#if !MEDIAPIPE_DISABLE_GPU
-
 #include "mediapipe/gpu/gpu_buffer.h"
 #include "mediapipe/gpu/gpu_buffer_format.h"
+#include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
+#include "mediapipe/gpu/image_frame_view.h"
+
+#if !MEDIAPIPE_DISABLE_GPU
 
 #if defined(__APPLE__)
 #include <CoreVideo/CoreVideo.h>
 
 #include "mediapipe/objc/CFHolder.h"
 #include "mediapipe/objc/util.h"
-#if !TARGET_OS_OSX  // iOS, use CVPixelBuffer.
-#define MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER 1
-#endif  // TARGET_OS_OSX
 #endif  // defined(__APPLE__)
 
 #if !MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER  // OSX, use GL textures.
@@ -73,15 +71,15 @@ class Image {
   // Creates an Image representing the same image content as the ImageFrame
   // the input shared pointer points to, and retaining shared ownership.
   explicit Image(ImageFrameSharedPtr image_frame)
-      : image_frame_(std::move(image_frame)) {
+      : gpu_buffer_(std::make_shared<GpuBufferStorageImageFrame>(
+            std::move(image_frame))) {
     use_gpu_ = false;
-    pixel_mutex_ = std::make_shared<absl::Mutex>();
   }
 
   // CPU getters.
-  const ImageFrameSharedPtr& GetImageFrameSharedPtr() const {
-    if (use_gpu_ == true) ConvertToCpu();
-    return image_frame_;
+  ImageFrameSharedPtr GetImageFrameSharedPtr() const {
+    // Write view currently because the return type does not point to const IF.
+    return gpu_buffer_.GetWriteView<ImageFrame>();
   }
 
   // Creates an Image representing the same image content as the input GPU
@@ -99,19 +97,18 @@ class Image {
   explicit Image(mediapipe::GpuBuffer gpu_buffer) {
     use_gpu_ = true;
     gpu_buffer_ = gpu_buffer;
-    pixel_mutex_ = std::make_shared<absl::Mutex>();
   }
 
   // GPU getters.
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   CVPixelBufferRef GetCVPixelBufferRef() const {
     if (use_gpu_ == false) ConvertToGpu();
-    return gpu_buffer_.GetCVPixelBufferRef();
+    return mediapipe::GetCVPixelBufferRef(gpu_buffer_);
   }
 #else
   mediapipe::GlTextureBufferSharedPtr GetGlTextureBufferSharedPtr() const {
     if (use_gpu_ == false) ConvertToGpu();
-    return gpu_buffer_.GetGlTextureBufferSharedPtr();
+    return gpu_buffer_.internal_storage<mediapipe::GlTextureBuffer>();
   }
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   // Get a GPU view. Automatically uploads from CPU if needed.
@@ -128,9 +125,7 @@ class Image {
   int step() const;  // Row size in bytes.
   bool UsesGpu() const { return use_gpu_; }
   ImageFormat::Format image_format() const;
-#if !MEDIAPIPE_DISABLE_GPU
   mediapipe::GpuBufferFormat format() const;
-#endif  // !MEDIAPIPE_DISABLE_GPU
 
   // Converts to true iff valid.
   explicit operator bool() const { return operator!=(nullptr); }
@@ -147,8 +142,8 @@ class Image {
 
   // Lock/Unlock pixel data.
   // Should be used exclusively by the PixelLock helper class.
-  void LockPixels() const ABSL_EXCLUSIVE_LOCK_FUNCTION(pixel_mutex_);
-  void UnlockPixels() const ABSL_UNLOCK_FUNCTION(pixel_mutex_);
+  void LockPixels() const ABSL_EXCLUSIVE_LOCK_FUNCTION();
+  void UnlockPixels() const ABSL_UNLOCK_FUNCTION();
 
   // Helper utility for GPU->CPU data transfer.
   bool ConvertToCpu() const;
@@ -157,75 +152,32 @@ class Image {
   bool ConvertToGpu() const;
 
  private:
-#if !MEDIAPIPE_DISABLE_GPU
   mutable mediapipe::GpuBuffer gpu_buffer_;
-#endif  // !MEDIAPIPE_DISABLE_GPU
-  mutable ImageFrameSharedPtr image_frame_;
   mutable bool use_gpu_ = false;
-  mutable std::shared_ptr<absl::Mutex> pixel_mutex_;  // ImageFrame only.
 };
 
-inline int Image::width() const {
-#if !MEDIAPIPE_DISABLE_GPU
-  if (use_gpu_)
-    return gpu_buffer_.width();
-  else
-#endif  // !MEDIAPIPE_DISABLE_GPU
-    return image_frame_->Width();
-}
+inline int Image::width() const { return gpu_buffer_.width(); }
 
-inline int Image::height() const {
-#if !MEDIAPIPE_DISABLE_GPU
-  if (use_gpu_)
-    return gpu_buffer_.height();
-  else
-#endif  // !MEDIAPIPE_DISABLE_GPU
-    return image_frame_->Height();
-}
+inline int Image::height() const { return gpu_buffer_.height(); }
 
 inline ImageFormat::Format Image::image_format() const {
-#if !MEDIAPIPE_DISABLE_GPU
-  if (use_gpu_)
-    return mediapipe::ImageFormatForGpuBufferFormat(gpu_buffer_.format());
-  else
-#endif  // !MEDIAPIPE_DISABLE_GPU
-    return image_frame_->Format();
+  return mediapipe::ImageFormatForGpuBufferFormat(gpu_buffer_.format());
 }
 
-#if !MEDIAPIPE_DISABLE_GPU
 inline mediapipe::GpuBufferFormat Image::format() const {
-  if (use_gpu_)
-    return gpu_buffer_.format();
-  else
-    return mediapipe::GpuBufferFormatForImageFormat(image_frame_->Format());
+  return gpu_buffer_.format();
 }
-#endif  // !MEDIAPIPE_DISABLE_GPU
 
 inline bool Image::operator==(std::nullptr_t other) const {
-#if !MEDIAPIPE_DISABLE_GPU
-  if (use_gpu_)
-    return gpu_buffer_ == other;
-  else
-#endif  // !MEDIAPIPE_DISABLE_GPU
-    return image_frame_ == other;
+  return gpu_buffer_ == other;
 }
 
 inline bool Image::operator==(const Image& other) const {
-#if !MEDIAPIPE_DISABLE_GPU
-  if (use_gpu_)
-    return gpu_buffer_ == other.gpu_buffer_;
-  else
-#endif  // !MEDIAPIPE_DISABLE_GPU
-    return image_frame_ == other.image_frame_;
+  return gpu_buffer_ == other.gpu_buffer_;
 }
 
 inline Image& Image::operator=(std::nullptr_t other) {
-#if !MEDIAPIPE_DISABLE_GPU
-  if (use_gpu_)
-    gpu_buffer_ = other;
-  else
-#endif  // !MEDIAPIPE_DISABLE_GPU
-    image_frame_ = other;
+  gpu_buffer_ = other;
   return *this;
 }
 
@@ -234,19 +186,14 @@ inline int Image::channels() const {
 }
 
 inline int Image::step() const {
-  if (use_gpu_)
-    return width() * channels() *
-           ImageFrame::ByteDepthForFormat(image_format());
-  else
-    return image_frame_->WidthStep();
+  return gpu_buffer_.GetReadView<ImageFrame>()->WidthStep();
 }
 
 inline void Image::LockPixels() const {
-  pixel_mutex_->Lock();
   ConvertToCpu();  // Download data if necessary.
 }
 
-inline void Image::UnlockPixels() const { pixel_mutex_->Unlock(); }
+inline void Image::UnlockPixels() const {}
 
 // Helper class for getting access to Image CPU data,
 // and handles automatically locking/unlocking CPU data access.
@@ -268,7 +215,10 @@ class PixelReadLock {
  public:
   explicit PixelReadLock(const Image& image) {
     buffer_ = &image;
-    if (buffer_) buffer_->LockPixels();
+    if (buffer_) {
+      buffer_->LockPixels();
+      frame_ = buffer_->GetImageFrameSharedPtr();
+    }
   }
   ~PixelReadLock() {
     if (buffer_) buffer_->UnlockPixels();
@@ -276,10 +226,7 @@ class PixelReadLock {
   PixelReadLock(const PixelReadLock&) = delete;
 
   const uint8* Pixels() const {
-    if (buffer_ && !buffer_->UsesGpu()) {
-      ImageFrame* frame = buffer_->GetImageFrameSharedPtr().get();
-      if (frame) return frame->PixelData();
-    }
+    if (frame_) return frame_->PixelData();
     return nullptr;
   }
 
@@ -287,13 +234,17 @@ class PixelReadLock {
 
  private:
   const Image* buffer_ = nullptr;
+  std::shared_ptr<ImageFrame> frame_;
 };
 
 class PixelWriteLock {
  public:
   explicit PixelWriteLock(Image* image) {
     buffer_ = image;
-    if (buffer_) buffer_->LockPixels();
+    if (buffer_) {
+      buffer_->LockPixels();
+      frame_ = buffer_->GetImageFrameSharedPtr();
+    }
   }
   ~PixelWriteLock() {
     if (buffer_) buffer_->UnlockPixels();
@@ -301,10 +252,7 @@ class PixelWriteLock {
   PixelWriteLock(const PixelWriteLock&) = delete;
 
   uint8* Pixels() {
-    if (buffer_ && !buffer_->UsesGpu()) {
-      ImageFrame* frame = buffer_->GetImageFrameSharedPtr().get();
-      if (frame) return frame->MutablePixelData();
-    }
+    if (frame_) return frame_->MutablePixelData();
     return nullptr;
   }
 
@@ -312,6 +260,7 @@ class PixelWriteLock {
 
  private:
   const Image* buffer_ = nullptr;
+  std::shared_ptr<ImageFrame> frame_;
 };
 
 }  // namespace mediapipe

mediapipe/framework/formats/image_opencv.cc

@@ -77,7 +77,16 @@ int GetMatType(const mediapipe::ImageFormat::Format format) {
 namespace mediapipe {
 namespace formats {
 
-cv::Mat MatView(const mediapipe::Image* image) {
+std::shared_ptr<cv::Mat> MatView(const mediapipe::Image* image) {
+  // Used to hold the lock through the Mat's lifetime.
+  struct MatWithPixelLock {
+    // Constructor needed because you cannot use aggregate initialization with
+    // std::make_shared.
+    MatWithPixelLock(mediapipe::Image* image) : lock(image) {}
+    mediapipe::PixelWriteLock lock;
+    cv::Mat mat;
+  };
+
   const int dims = 2;
   const int sizes[] = {image->height(), image->width()};
   const int type =
@@ -85,18 +94,22 @@ cv::Mat MatView(const mediapipe::Image* image) {
   const size_t steps[] = {static_cast<size_t>(image->step()),
                           static_cast<size_t>(ImageFrame::ByteDepthForFormat(
                               image->image_format()))};
-  mediapipe::PixelWriteLock dst_lock(const_cast<mediapipe::Image*>(image));
-  uint8* data_ptr = dst_lock.Pixels();
+  auto owner =
+      std::make_shared<MatWithPixelLock>(const_cast<mediapipe::Image*>(image));
+  uint8* data_ptr = owner->lock.Pixels();
   CHECK(data_ptr != nullptr);
   // Use Image to initialize in-place. Image still owns memory.
   if (steps[0] == sizes[1] * image->channels() *
                       ImageFrame::ByteDepthForFormat(image->image_format())) {
     // Contiguous memory optimization. See b/78570764
-    return cv::Mat(dims, sizes, type, data_ptr);
+    owner->mat = cv::Mat(dims, sizes, type, data_ptr);
   } else {
     // Custom width step.
-    return cv::Mat(dims, sizes, type, data_ptr, steps);
+    owner->mat = cv::Mat(dims, sizes, type, data_ptr, steps);
   }
+  // Aliasing constructor makes a shared_ptr<Mat> which keeps the whole
+  // MatWithPixelLock alive.
+  return std::shared_ptr<cv::Mat>(owner, &owner->mat);
 }
 }  // namespace formats
 }  // namespace mediapipe

mediapipe/framework/formats/image_opencv.h

@@ -29,7 +29,9 @@ namespace formats {
 // the const modifier is lost.  The caller must be careful
 // not to use the returned object to modify the data in a const Image,
 // even though the returned data is mutable.
-cv::Mat MatView(const mediapipe::Image* image);
+// Note: this returns a shared_ptr so it can keep the CPU memory referenced
+// by the Mat alive.
+std::shared_ptr<cv::Mat> MatView(const mediapipe::Image* image);
 
 }  // namespace formats
 }  // namespace mediapipe

mediapipe/framework/formats/matrix.h

@@ -39,7 +39,7 @@ void MatrixDataProtoFromMatrix(const Matrix& matrix, MatrixData* matrix_data);
 void MatrixFromMatrixDataProto(const MatrixData& matrix_data, Matrix* matrix);
 
 #if !defined(MEDIAPIPE_MOBILE) && !defined(MEDIAPIPE_LITE)
-// Produce a Text format MatrixData std::string.  Mainly useful for test code.
+// Produce a Text format MatrixData string.  Mainly useful for test code.
 std::string MatrixAsTextProto(const Matrix& matrix);
 // Produce a Matrix from a text format MatrixData proto representation.
 void MatrixFromTextProto(const std::string& text_proto, Matrix* matrix);

mediapipe/framework/formats/tensor.h

@@ -76,7 +76,7 @@ class Tensor {
 
  public:
   // No resources are allocated here.
-  enum class ElementType { kNone, kFloat16, kFloat32 };
+  enum class ElementType { kNone, kFloat16, kFloat32, kUInt8 };
   struct Shape {
     Shape() = default;
     Shape(std::initializer_list<int> dimensions) : dims(dimensions) {}
@@ -215,6 +215,8 @@ class Tensor {
         return 2;
       case ElementType::kFloat32:
         return sizeof(float);
+      case ElementType::kUInt8:
+        return 1;
     }
   }
   int bytes() const { return shape_.num_elements() * element_size(); }
@@ -278,6 +280,11 @@ class Tensor {
 #endif  // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
 };
 
+int BhwcBatchFromShape(const Tensor::Shape& shape);
+int BhwcHeightFromShape(const Tensor::Shape& shape);
+int BhwcWidthFromShape(const Tensor::Shape& shape);
+int BhwcDepthFromShape(const Tensor::Shape& shape);
+
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_H_

mediapipe/framework/formats/tensor_internal.h

@@ -16,14 +16,22 @@
 #define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_INTERNAL_H_
 
 #include <cstdint>
+#include <type_traits>
+
+#include "mediapipe/framework/tool/type_util.h"
 
 namespace mediapipe {
 
 // Generates unique view id at compile-time using FILE and LINE.
-#define TENSOR_UNIQUE_VIEW_TYPE_ID()                          \
-  static constexpr uint64_t kId = tensor_internal::FnvHash64( \
+#define TENSOR_UNIQUE_VIEW_TYPE_ID()                       \
+  static inline uint64_t kId = tensor_internal::FnvHash64( \
       __FILE__, tensor_internal::FnvHash64(TENSOR_INT_TO_STRING(__LINE__)))
 
+// Generates unique view id at compile-time using FILE and LINE and Type of the
+// template view's argument.
+#define TENSOR_UNIQUE_VIEW_TYPE_ID_T(T) \
+  static inline uint64_t kId = tool::GetTypeHash<T>();
+
 namespace tensor_internal {
 
 #define TENSOR_INT_TO_STRING2(x) #x
@@ -36,6 +44,21 @@ constexpr uint64_t kFnvOffsetBias = 0xcbf29ce484222325;
 constexpr uint64_t FnvHash64(const char* str, uint64_t hash = kFnvOffsetBias) {
   return (str[0] == 0) ? hash : FnvHash64(str + 1, (hash ^ str[0]) * kFnvPrime);
 }
+
+template <typename... Ts>
+struct TypeList {
+  static constexpr std::size_t size{sizeof...(Ts)};
+};
+template <typename, typename>
+struct TypeInList {};
+template <typename T, typename... Ts>
+struct TypeInList<T, TypeList<T, Ts...>>
+    : std::integral_constant<std::size_t, 0> {};
+template <typename T, typename TOther, typename... Ts>
+struct TypeInList<T, TypeList<TOther, Ts...>>
+    : std::integral_constant<std::size_t,
+                             1 + TypeInList<T, TypeList<Ts...>>::value> {};
+
 }  // namespace tensor_internal
 }  // namespace mediapipe
 

mediapipe/framework/graph_validation_test.cc

@@ -355,8 +355,8 @@ TEST(GraphValidationTest, OptionalInputNotProvidedForSubgraphCalculator) {
     output_stream: "OUTPUT:output_0"
     node {
       calculator: "OptionalSideInputTestCalculator"
-      input_side_packet: "SIDEINPUT:input_0"  # std::string
-      output_stream: "OUTPUT:output_0"        # std::string
+      input_side_packet: "SIDEINPUT:input_0"  # string
+      output_stream: "OUTPUT:output_0"        # string
     }
   )pb");
 
@@ -366,7 +366,7 @@ TEST(GraphValidationTest, OptionalInputNotProvidedForSubgraphCalculator) {
     output_stream: "OUTPUT:foo_out"
     node {
       calculator: "PassThroughGraph"
-      output_stream: "OUTPUT:foo_out"  # std::string
+      output_stream: "OUTPUT:foo_out"  # string
     }
   )pb");
 
@@ -406,10 +406,10 @@ TEST(GraphValidationTest, MultipleOptionalInputsForSubgraph) {
     output_stream: "OUTPUT:output_0"
     node {
       calculator: "OptionalSideInputTestCalculator"
-      input_side_packet: "SIDEINPUT:input_0"  # std::string
+      input_side_packet: "SIDEINPUT:input_0"  # string
       input_stream: "SELECT:select"
       input_stream: "ENABLE:enable"
-      output_stream: "OUTPUT:output_0"  # std::string
+      output_stream: "OUTPUT:output_0"  # string
     }
   )pb");
 
@@ -421,7 +421,7 @@ TEST(GraphValidationTest, MultipleOptionalInputsForSubgraph) {
     node {
       calculator: "PassThroughGraph"
       input_stream: "SELECT:foo_select"
-      output_stream: "OUTPUT:foo_out"  # std::string
+      output_stream: "OUTPUT:foo_out"  # string
     }
   )pb");
 

mediapipe/framework/input_stream_handler.h

@@ -147,8 +147,7 @@ class InputStreamHandler {
 
   void Close();
 
-  // Returns a std::string that concatenates the stream names of all managed
-  // streams.
+  // Returns a string that concatenates the stream names of all managed streams.
   std::string DebugStreamNames() const;
 
   // Keeps scheduling new invocations until 1) the node is not ready or 2) the

mediapipe/framework/input_stream_shard.h

@@ -51,7 +51,7 @@ class InputStreamShard : public InputStream {
     return !packet_queue_.empty() ? packet_queue_.front() : empty_packet_;
   }
 
-  // Returns a reference to the name std::string of the InputStreamManager.
+  // Returns a reference to the name string of the InputStreamManager.
   const std::string& Name() const { return *name_; }
 
   bool IsDone() const override { return is_done_; }
@@ -75,7 +75,7 @@ class InputStreamShard : public InputStream {
   std::queue<Packet> packet_queue_;
   Packet empty_packet_;
 
-  // Pointer to the name std::string of the InputStreamManager.
+  // Pointer to the name string of the InputStreamManager.
   const std::string* name_;
   bool is_done_;
 

mediapipe/framework/lifetime_tracker.h

@@ -18,6 +18,7 @@
 #include <atomic>
 
 #include "absl/memory/memory.h"
+#include "absl/synchronization/mutex.h"
 
 namespace mediapipe {
 
@@ -33,9 +34,13 @@ class LifetimeTracker {
   class Object {
    public:
     explicit Object(LifetimeTracker* tracker) : tracker_(tracker) {
+      absl::MutexLock lock(&tracker_->mutex_);
       ++tracker_->live_count_;
     }
-    ~Object() { --tracker_->live_count_; }
+    ~Object() {
+      absl::MutexLock lock(&tracker_->mutex_);
+      --tracker_->live_count_;
+    }
 
    private:
     LifetimeTracker* const tracker_;
@@ -47,10 +52,26 @@ class LifetimeTracker {
   }
 
   // Returns the number of tracked objects currently alive.
-  int live_count() { return live_count_; }
+  int live_count() {
+    absl::MutexLock lock(&mutex_);
+    return live_count_;
+  }
+
+  // Waits for all instances of Object to be destroyed / live_count to reach
+  // zero. Returns true if this occurred within the timeout, false otherwise.
+  bool WaitForAllObjectsToDie(
+      absl::Duration timeout = absl::InfiniteDuration()) {
+    // Condition takes a function pointer. Prefixing the lambda with a +
+    // resolves it to a pointer.
+    absl::Condition check_count(
+        +[](int* value) { return *value == 0; }, &live_count_);
+    absl::MutexLock lock(&mutex_);
+    return mutex_.AwaitWithTimeout(check_count, timeout);
+  }
 
  private:
-  std::atomic<int> live_count_ = ATOMIC_VAR_INIT(0);
+  absl::Mutex mutex_;
+  int live_count_ ABSL_GUARDED_BY(mutex_) = 0;
 };
 
 }  // namespace mediapipe

mediapipe/framework/mediapipe_cc_test.bzl

@@ -10,10 +10,17 @@ def mediapipe_cc_test(
         size = None,
         tags = [],
         timeout = None,
+        args = [],
         additional_deps = DEFAULT_ADDITIONAL_TEST_DEPS,
+        # ios_unit_test arguments
+        ios_minimum_os_version = "9.0",
+        # android_cc_test arguments
+        open_gl_driver = None,
+        emulator_mini_boot = True,
+        requires_full_emulation = True,
+        # wasm_web_test arguments
+        browsers = None,
         **kwargs):
-    # Note: additional_deps are MediaPipe-specific test support deps added by default.
-    # They are provided as a default argument so they can be disabled if desired.
     native.cc_library(
         name = name + "_lib",
         testonly = 1,

mediapipe/framework/more_selects.bzl

@@ -0,0 +1,76 @@
+"""More utilities to help with selects."""
+
+load("@bazel_skylib//lib:selects.bzl", "selects")
+
+# From selects.bzl, but it's not public there.
+def _config_setting_always_true(name, visibility):
+    """Returns a config_setting with the given name that's always true.
+
+    This is achieved by constructing a two-entry OR chain where each
+    config_setting takes opposite values of a boolean flag.
+    """
+    name_on = name + "_stamp_binary_on_check"
+    name_off = name + "_stamp_binary_off_check"
+    native.config_setting(
+        name = name_on,
+        values = {"stamp": "1"},
+    )
+    native.config_setting(
+        name = name_off,
+        values = {"stamp": "0"},
+    )
+    return selects.config_setting_group(
+        name = name,
+        visibility = visibility,
+        match_any = [
+            ":" + name_on,
+            ":" + name_off,
+        ],
+    )
+
+def _config_setting_always_false(name, visibility):
+    """Returns a config_setting with the given name that's always false.
+
+    This is achieved by constructing a two-entry AND chain where each
+    config_setting takes opposite values of a boolean flag.
+    """
+    name_on = name + "_stamp_binary_on_check"
+    name_off = name + "_stamp_binary_off_check"
+    native.config_setting(
+        name = name_on,
+        values = {"stamp": "1"},
+    )
+    native.config_setting(
+        name = name_off,
+        values = {"stamp": "0"},
+    )
+    return selects.config_setting_group(
+        name = name,
+        visibility = visibility,
+        match_all = [
+            ":" + name_on,
+            ":" + name_off,
+        ],
+    )
+
+def _config_setting_negation(name, negate, visibility = None):
+    _config_setting_always_true(
+        name = name + "_true",
+        visibility = visibility,
+    )
+    _config_setting_always_false(
+        name = name + "_false",
+        visibility = visibility,
+    )
+    native.alias(
+        name = name,
+        actual = select({
+            "//conditions:default": ":%s_true" % name,
+            negate: ":%s_false" % name,
+        }),
+        visibility = visibility,
+    )
+
+more_selects = struct(
+    config_setting_negation = _config_setting_negation,
+)

mediapipe/framework/packet.cc

@@ -113,7 +113,8 @@ absl::Status Packet::ValidateAsType(const tool::TypeInfo& type_info) const {
                      MediaPipeTypeStringOrDemangled(type_info),
                      ", but received an empty Packet."));
   }
-  bool holder_is_right_type = holder_->GetTypeId() == type_info.hash_code();
+  bool holder_is_right_type =
+      holder_->GetTypeInfo().hash_code() == type_info.hash_code();
   if (ABSL_PREDICT_FALSE(!holder_is_right_type)) {
     return absl::InvalidArgumentError(absl::StrCat(
         "The Packet stores \"", holder_->DebugTypeName(), "\", but \"",

mediapipe/framework/packet.h

@@ -189,7 +189,11 @@ class Packet {
 
   // Get the type id for the underlying type stored in the Packet.
   // Crashes if IsEmpty() == true.
-  size_t GetTypeId() const;
+  size_t GetTypeId() const { return GetTypeInfo().hash_code(); }
+
+  // Get the type info for the underlying type stored in the Packet.
+  // Crashes if IsEmpty() == true.
+  const tool::TypeInfo& GetTypeInfo() const;
 
   // Returns the timestamp.
   class Timestamp Timestamp() const;
@@ -201,9 +205,9 @@ class Packet {
 
   // Returns the type name.  If the packet is empty or the type is not
   // registered (with MEDIAPIPE_REGISTER_TYPE or companion macros) then
-  // the empty std::string is returned.
+  // the empty string is returned.
   std::string RegisteredTypeName() const;
-  // Returns a std::string with the best guess at the type name.
+  // Returns a string with the best guess at the type name.
   std::string DebugTypeName() const;
 
  private:
@@ -220,6 +224,7 @@ class Packet {
   friend std::shared_ptr<packet_internal::HolderBase>
   packet_internal::GetHolderShared(Packet&& packet);
 
+  friend class PacketType;
   absl::Status ValidateAsType(const tool::TypeInfo& type_info) const;
 
   std::shared_ptr<packet_internal::HolderBase> holder_;
@@ -364,15 +369,15 @@ class HolderBase {
   virtual ~HolderBase();
   template <typename T>
   bool PayloadIsOfType() const {
-    return GetTypeId() == tool::GetTypeHash<T>();
+    return GetTypeInfo().hash_code() == tool::GetTypeHash<T>();
   }
-  // Returns a printable std::string identifying the type stored in the holder.
+  // Returns a printable string identifying the type stored in the holder.
   virtual const std::string DebugTypeName() const = 0;
   // Returns the registered type name if it's available, otherwise the
-  // empty std::string.
+  // empty string.
   virtual const std::string RegisteredTypeName() const = 0;
   // Get the type id of the underlying data type.
-  virtual size_t GetTypeId() const = 0;
+  virtual const tool::TypeInfo& GetTypeInfo() const = 0;
   // Downcasts this to Holder<T>.  Returns nullptr if deserialization
   // failed or if the requested type is not what is stored.
   template <typename T>
@@ -440,7 +445,7 @@ ConvertToVectorOfProtoMessageLitePtrs(const T* data,
 }
 
 // This registry is used to create Holders of the right concrete C++ type given
-// a proto type std::string (which is used as the registration key).
+// a proto type string (which is used as the registration key).
 class MessageHolderRegistry
     : public GlobalFactoryRegistry<std::unique_ptr<HolderBase>> {};
 
@@ -505,7 +510,7 @@ class Holder : public HolderBase {
     HolderSupport<T>::EnsureStaticInit();
     return *ptr_;
   }
-  size_t GetTypeId() const final { return tool::GetTypeHash<T>(); }
+  const tool::TypeInfo& GetTypeInfo() const final { return tool::TypeId<T>(); }
   // Releases the underlying data pointer and transfers the ownership to a
   // unique pointer.
   // This method is dangerous and is only used by Packet::Consume() if the
@@ -741,9 +746,9 @@ inline Packet& Packet::operator=(Packet&& packet) {
 
 inline bool Packet::IsEmpty() const { return holder_ == nullptr; }
 
-inline size_t Packet::GetTypeId() const {
+inline const tool::TypeInfo& Packet::GetTypeInfo() const {
   CHECK(holder_);
-  return holder_->GetTypeId();
+  return holder_->GetTypeInfo();
 }
 
 template <typename T>

mediapipe/framework/packet_type.cc

@@ -19,57 +19,55 @@
 #include <unordered_set>
 #include <utility>
 
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
+#include "absl/types/span.h"
+#include "absl/types/variant.h"
 #include "mediapipe/framework/port/canonical_errors.h"
 #include "mediapipe/framework/port/logging.h"
 #include "mediapipe/framework/port/map_util.h"
 #include "mediapipe/framework/port/source_location.h"
 #include "mediapipe/framework/port/status_builder.h"
 #include "mediapipe/framework/tool/status_util.h"
+#include "mediapipe/framework/tool/type_util.h"
 #include "mediapipe/framework/tool/validate_name.h"
+#include "mediapipe/framework/type_map.h"
 
 namespace mediapipe {
 
-PacketType::PacketType()
-    : initialized_(false),
-      no_packets_allowed_(true),
-      validate_method_(nullptr),
-      type_name_("[Undefined Type]"),
-      same_as_(nullptr) {}
+absl::Status PacketType::AcceptAny(const TypeSpec& type) {
+  return absl::OkStatus();
+}
+
+absl::Status PacketType::AcceptNone(const TypeSpec& type) {
+  auto* special = absl::get_if<SpecialType>(&type);
+  if (special &&
+      (special->accept_fn_ == AcceptNone || special->accept_fn_ == AcceptAny))
+    return absl::OkStatus();
+  return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+         << "No packets are allowed for type: [No Type]";
+}
 
 PacketType& PacketType::SetAny() {
-  no_packets_allowed_ = false;
-  validate_method_ = nullptr;
-  same_as_ = nullptr;
-  type_name_ = "[Any Type]";
-  initialized_ = true;
+  type_spec_ = SpecialType{"[Any Type]", &AcceptAny};
   return *this;
 }
 
 PacketType& PacketType::SetNone() {
-  no_packets_allowed_ = true;
-  validate_method_ = nullptr;
-  same_as_ = nullptr;
-  type_name_ = "[No Type]";
-  initialized_ = true;
+  type_spec_ = SpecialType{"[No Type]", &AcceptNone};
   return *this;
 }
 
 PacketType& PacketType::SetSameAs(const PacketType* type) {
   // TODO Union sets together when SetSameAs is called multiple times.
-  no_packets_allowed_ = false;
-  validate_method_ = nullptr;
-  same_as_ = type->GetSameAs();
-  type_name_ = "";
-
-  if (same_as_ == this) {
+  auto same_as = type->GetSameAs();
+  if (same_as == this) {
     // We're the root of the union-find tree.  There's a cycle, which
     // means we might as well be an "Any" type.
-    same_as_ = nullptr;
+    return SetAny();
   }
-
-  initialized_ = true;
+  type_spec_ = SameAs{same_as};
   return *this;
 }
 
@@ -78,10 +76,19 @@ PacketType& PacketType::Optional() {
   return *this;
 }
 
-bool PacketType::IsInitialized() const { return initialized_; }
+bool PacketType::IsInitialized() const {
+  return !absl::holds_alternative<absl::monostate>(type_spec_);
+}
+
+const PacketType* PacketType::SameAsPtr() const {
+  auto* same_as = absl::get_if<SameAs>(&type_spec_);
+  if (same_as) return same_as->other;
+  return nullptr;
+}
 
 PacketType* PacketType::GetSameAs() {
-  if (!same_as_) {
+  auto* same_as = SameAsPtr();
+  if (!same_as) {
     return this;
   }
   // Don't optimize the union-find algorithm, since updating the pointer
@@ -91,89 +98,174 @@ PacketType* PacketType::GetSameAs() {
   // make the current set point to the root of the other tree.
   // TODO Remove const_cast by making SetSameAs take a non-const
   // PacketType*.
-  return const_cast<PacketType*>(same_as_->GetSameAs());
+  return const_cast<PacketType*>(same_as->GetSameAs());
 }
 
 const PacketType* PacketType::GetSameAs() const {
-  if (!same_as_) {
+  auto* same_as = SameAsPtr();
+  if (!same_as) {
     return this;
   }
   // See comments in non-const variant.
-  return same_as_->GetSameAs();
+  return same_as->GetSameAs();
 }
 
 bool PacketType::IsAny() const {
-  return !no_packets_allowed_ && validate_method_ == nullptr &&
-         same_as_ == nullptr;
+  auto* special = absl::get_if<SpecialType>(&type_spec_);
+  return special && special->accept_fn_ == AcceptAny;
 }
 
-bool PacketType::IsNone() const { return no_packets_allowed_; }
+bool PacketType::IsNone() const {
+  auto* special = absl::get_if<SpecialType>(&type_spec_);
+  // The tests currently require that an uninitialized PacketType return true
+  // for IsNone. TODO: change it?
+  return !IsInitialized() || (special && special->accept_fn_ == AcceptNone);
+}
+
+bool PacketType::IsOneOf() const {
+  return absl::holds_alternative<MultiType>(type_spec_);
+}
+
+bool PacketType::IsExactType() const {
+  return absl::holds_alternative<const tool::TypeInfo*>(type_spec_);
+}
 
 const std::string* PacketType::RegisteredTypeName() const {
-  if (same_as_) {
-    return GetSameAs()->RegisteredTypeName();
-  }
-  return registered_type_name_ptr_;
+  if (auto* same_as = SameAsPtr()) return same_as->RegisteredTypeName();
+  if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec_))
+    return MediaPipeTypeStringFromTypeId((**type_info).hash_code());
+  if (auto* multi_type = absl::get_if<MultiType>(&type_spec_))
+    return multi_type->registered_type_name;
+  return nullptr;
 }
 
-const std::string PacketType::DebugTypeName() const {
-  if (same_as_) {
+namespace internal {
+
+struct TypeInfoFormatter {
+  void operator()(std::string* out, const tool::TypeInfo& t) const {
+    absl::StrAppend(out, MediaPipeTypeStringOrDemangled(t));
+  }
+};
+
+template <class Formatter>
+class QuoteFormatter {
+ public:
+  explicit QuoteFormatter(Formatter&& f) : f_(std::forward<Formatter>(f)) {}
+
+  template <typename T>
+  void operator()(std::string* out, const T& t) const {
+    absl::StrAppend(out, "\"");
+    f_(out, t);
+    absl::StrAppend(out, "\"");
+  }
+
+ private:
+  Formatter f_;
+};
+template <class Formatter>
+explicit QuoteFormatter(Formatter f) -> QuoteFormatter<Formatter>;
+
+}  // namespace internal
+
+std::string PacketType::TypeNameForOneOf(TypeInfoSpan types) {
+  return absl::StrCat(
+      "OneOf<",
+      absl::StrJoin(types, ", ",
+                    absl::DereferenceFormatter(internal::TypeInfoFormatter())),
+      ">");
+}
+
+std::string PacketType::DebugTypeName() const {
+  if (auto* same_as = absl::get_if<SameAs>(&type_spec_)) {
     // Construct a name based on the current chain of same_as_ links
     // (which may change when the framework expands out Any-type).
-    return absl::StrCat("[Same Type As ", GetSameAs()->DebugTypeName(), "]");
+    return absl::StrCat("[Same Type As ",
+                        same_as->other->GetSameAs()->DebugTypeName(), "]");
   }
-  return type_name_;
+  if (auto* special = absl::get_if<SpecialType>(&type_spec_)) {
+    return special->name_;
+  }
+  if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec_)) {
+    return MediaPipeTypeStringOrDemangled(**type_info);
+  }
+  if (auto* multi_type = absl::get_if<MultiType>(&type_spec_)) {
+    return TypeNameForOneOf(multi_type->types);
+  }
+  return "[Undefined Type]";
+}
+
+static bool HaveCommonType(absl::Span<const tool::TypeInfo* const> types1,
+                           absl::Span<const tool::TypeInfo* const> types2) {
+  for (const auto& first : types1) {
+    for (const auto& second : types2) {
+      if (first->hash_code() == second->hash_code()) {
+        return true;
+      }
+    }
+  }
+  return false;
 }
 
 absl::Status PacketType::Validate(const Packet& packet) const {
-  if (!initialized_) {
+  if (!IsInitialized()) {
     return absl::InvalidArgumentError(
         "Uninitialized PacketType was used for validation.");
   }
-  if (same_as_) {
+  if (SameAsPtr()) {
     // Cycles are impossible at this stage due to being checked for
     // in SetSameAs().
     return GetSameAs()->Validate(packet);
   }
-  if (no_packets_allowed_) {
-    return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
-           << "No packets are allowed for type: " << type_name_;
+  if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec_)) {
+    return packet.ValidateAsType(**type_info);
   }
-  if (validate_method_ != nullptr) {
-    return (packet.*validate_method_)();
-  }
-  // The PacketType is the Any Type.
   if (packet.IsEmpty()) {
     return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
-           << "Empty packets are not allowed for type: " << type_name_;
+           << "Empty packets are not allowed for type: " << DebugTypeName();
+  }
+  if (auto* multi_type = absl::get_if<MultiType>(&type_spec_)) {
+    auto* packet_type = &packet.GetTypeInfo();
+    if (HaveCommonType(multi_type->types, absl::MakeSpan(&packet_type, 1))) {
+      return absl::OkStatus();
+    } else {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "The Packet stores \"", packet.DebugTypeName(), "\", but one of ",
+          absl::StrJoin(multi_type->types, ", ",
+                        absl::DereferenceFormatter(internal::QuoteFormatter(
+                            internal::TypeInfoFormatter()))),
+          " was requested."));
+    }
+  }
+  if (auto* special = absl::get_if<SpecialType>(&type_spec_)) {
+    return special->accept_fn_(&packet.GetTypeInfo());
   }
   return absl::OkStatus();
 }
 
+PacketType::TypeInfoSpan PacketType::GetTypeSpan(const TypeSpec& type_spec) {
+  if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec))
+    return absl::MakeSpan(type_info, 1);
+  if (auto* multi_type = absl::get_if<MultiType>(&type_spec))
+    return multi_type->types;
+  return {};
+}
+
 bool PacketType::IsConsistentWith(const PacketType& other) const {
   const PacketType* type1 = GetSameAs();
   const PacketType* type2 = other.GetSameAs();
 
-  if (type1->validate_method_ == nullptr ||
-      type2->validate_method_ == nullptr) {
-    // type1 or type2 either accepts anything or nothing.
-    if (type1->validate_method_ == nullptr && !type1->no_packets_allowed_) {
-      // type1 accepts anything.
-      return true;
-    }
-    if (type2->validate_method_ == nullptr && !type2->no_packets_allowed_) {
-      // type2 accepts anything.
-      return true;
-    }
-    if (type1->no_packets_allowed_ && type2->no_packets_allowed_) {
-      // type1 and type2 both accept nothing.
-      return true;
-    }
-    // The only special case left is that only one of "type1" or "type2"
-    // accepts nothing, which means there is no match.
-    return false;
+  TypeInfoSpan types1 = GetTypeSpan(type1->type_spec_);
+  TypeInfoSpan types2 = GetTypeSpan(type2->type_spec_);
+  if (!types1.empty() && !types2.empty()) {
+    return HaveCommonType(types1, types2);
   }
-  return type1->validate_method_ == type2->validate_method_;
+  if (auto* special1 = absl::get_if<SpecialType>(&type1->type_spec_)) {
+    return special1->accept_fn_(type2->type_spec_).ok();
+  }
+  if (auto* special2 = absl::get_if<SpecialType>(&type2->type_spec_)) {
+    return special2->accept_fn_(type1->type_spec_).ok();
+  }
+  return false;
 }
 
 absl::Status ValidatePacketTypeSet(const PacketTypeSet& packet_type_set) {

mediapipe/framework/packet_type.h

@@ -23,12 +23,16 @@
 #include <vector>
 
 #include "absl/base/macros.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "mediapipe/framework/collection.h"
+#include "mediapipe/framework/deps/no_destructor.h"
 #include "mediapipe/framework/packet.h"
 #include "mediapipe/framework/packet_set.h"
 #include "mediapipe/framework/port/status.h"
+#include "mediapipe/framework/tool/type_util.h"
 #include "mediapipe/framework/tool/validate_name.h"
 #include "mediapipe/framework/type_map.h"
 
@@ -41,7 +45,7 @@ namespace mediapipe {
 class PacketType {
  public:
   // Creates an uninitialized PacketType.
-  PacketType();
+  PacketType() = default;
 
   // PacketType can be passed by value.
   PacketType(const PacketType&) = default;
@@ -63,6 +67,9 @@ class PacketType {
   // Specifically, using SetAny() still means that the stream has a type
   // but this particular calculator just doesn't care what it is.
   PacketType& SetAny();
+  // Sets the packet type to accept any of the provided types.
+  template <typename... T>
+  PacketType& SetOneOf();
   // Sets the packet type to not accept any packets.
   PacketType& SetNone();
   // Sets the PacketType to be the same as type.  This actually stores
@@ -80,6 +87,11 @@ class PacketType {
   bool IsAny() const;
   // Returns true if this PacketType allows nothing.
   bool IsNone() const;
+  // Returns true if this PacketType allows a set of types.
+  bool IsOneOf() const;
+  // Returns true if this PacketType allows one specific type.
+  bool IsExactType() const;
+  // Returns true if this port has been marked as optional.
   bool IsOptional() const { return optional_; }
 
   // Returns true iff this and other are consistent, meaning they do
@@ -101,26 +113,38 @@ class PacketType {
   const std::string* RegisteredTypeName() const;
   // Returns the type name.  Do not use this for validation, use
   // Validate() instead.
-  const std::string DebugTypeName() const;
+  std::string DebugTypeName() const;
 
  private:
-  // Typedef for the ValidateAsType() method in Packet that is used for
-  // type validation and identification.
-  typedef absl::Status (Packet::*ValidateMethodType)() const;
+  struct SameAs {
+    // This PacketType is the same as other.
+    // We don't do union-find optimizations in order to avoid a mutex.
+    const PacketType* other;
+  };
+  using TypeInfoSpan = absl::Span<const tool::TypeInfo* const>;
+  struct MultiType {
+    TypeInfoSpan types;
+    // TODO: refactor RegisteredTypeName, remove.
+    const std::string* registered_type_name;
+  };
+  struct SpecialType;
+  using TypeSpec = absl::variant<absl::monostate, const tool::TypeInfo*,
+                                 MultiType, SameAs, SpecialType>;
+  typedef absl::Status (*AcceptsTypeFn)(const TypeSpec& type);
+  struct SpecialType {
+    std::string name_;
+    AcceptsTypeFn accept_fn_;
+  };
+
+  static absl::Status AcceptAny(const TypeSpec& type);
+  static absl::Status AcceptNone(const TypeSpec& type);
+
+  const PacketType* SameAsPtr() const;
+  static TypeInfoSpan GetTypeSpan(const TypeSpec& type_spec);
+  static std::string TypeNameForOneOf(TypeInfoSpan types);
+
+  TypeSpec type_spec_;
 
-  // Records whether the packet type was set in any way.
-  bool initialized_;
-  // Don't allow any packets through.
-  bool no_packets_allowed_;
-  // Pointer to Packet::ValidateAsType<T>.
-  ValidateMethodType validate_method_;
-  // Type name as std::string.
-  std::string type_name_;
-  // The Registered type name or nullptr if the type isn't registered.
-  const std::string* registered_type_name_ptr_ = nullptr;
-  // If this is non-null then this PacketType is the same as same_as_.
-  // We don't do union-find optimizations in order to avoid a mutex.
-  const PacketType* same_as_;
   // Whether the corresponding port is optional.
   bool optional_ = false;
 };
@@ -164,7 +188,7 @@ class PacketTypeSetErrorHandler {
       for (const auto& entry : missing_->entries) {
         // Optional entries that were missing are not considered errors.
         if (!entry.second.IsOptional()) {
-          // Split them to keep the error std::string unchanged.
+          // Split them to keep the error string unchanged.
           std::pair<std::string, std::string> tag_idx =
               absl::StrSplit(entry.first, ':');
           missing_->errors.push_back(absl::StrCat("Failed to get tag \"",
@@ -235,12 +259,16 @@ absl::Status ValidatePacketTypeSet(const PacketTypeSet& packet_type_set);
 
 template <typename T>
 PacketType& PacketType::Set() {
-  initialized_ = true;
-  no_packets_allowed_ = false;
-  validate_method_ = &Packet::ValidateAsType<T>;
-  type_name_ = MediaPipeTypeStringOrDemangled<T>();
-  registered_type_name_ptr_ = MediaPipeTypeString<T>();
-  same_as_ = nullptr;
+  type_spec_ = &tool::TypeId<T>();
+  return *this;
+}
+
+template <typename... T>
+PacketType& PacketType::SetOneOf() {
+  static const NoDestructor<std::vector<const tool::TypeInfo*>> types{
+      {&tool::TypeId<T>()...}};
+  static const NoDestructor<std::string> name{TypeNameForOneOf(*types)};
+  type_spec_ = MultiType{*types, &*name};
   return *this;
 }
 

mediapipe/framework/port.h

@@ -60,18 +60,23 @@
 #define MEDIAPIPE_OPENGL_ES_30 300
 #define MEDIAPIPE_OPENGL_ES_31 310
 
+// NOTE: MEDIAPIPE_OPENGL_ES_VERSION macro represents the maximum OpenGL ES
+// version to build for. Runtime availability is _not_ guaranteed; in
+// particular, uses of OpenGL ES 3.1 should be guarded by a runtime check.
+// TODO: identify and fix code where macro is used incorrectly.
 #if MEDIAPIPE_DISABLE_GPU
 #define MEDIAPIPE_OPENGL_ES_VERSION 0
 #define MEDIAPIPE_METAL_ENABLED 0
 #else
 #if defined(MEDIAPIPE_ANDROID)
 #if defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
-#define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_20
+#define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_30
 #else
 #define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_31
 #endif
 #define MEDIAPIPE_METAL_ENABLED 0
 #elif defined(MEDIAPIPE_IOS)
+// TODO: use MEDIAPIPE_OPENGL_ES_30 for iOS as max version.
 #define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_20
 #define MEDIAPIPE_METAL_ENABLED 1
 #elif defined(MEDIAPIPE_OSX)

mediapipe/framework/profiler/BUILD

@@ -89,9 +89,11 @@ cc_library(
 cc_library(
     name = "graph_profiler_real",
     srcs = [
-        "gl_context_profiler.cc",
         "graph_profiler.cc",
-    ],
+    ] + select({
+        "//conditions:default": ["gl_context_profiler.cc"],
+        "//mediapipe/gpu:disable_gpu": [],
+    }),
     hdrs = [
         "graph_profiler.h",
     ],
@@ -100,31 +102,37 @@ cc_library(
     ],
     visibility = ["//visibility:private"],
     deps = [
-        ":graph_tracer",
         ":profiler_resource_util",
-        ":sharded_map",
+        ":graph_tracer",
         ":trace_buffer",
+        ":sharded_map",
         "//mediapipe/framework:calculator_cc_proto",
-        "//mediapipe/framework:calculator_context",
         "//mediapipe/framework:calculator_profile_cc_proto",
+        "//mediapipe/framework/port:integral_types",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+        "//mediapipe/framework/deps:clock",
+        "//mediapipe/framework:calculator_context",
         "//mediapipe/framework:executor",
         "//mediapipe/framework:validated_graph_config",
-        "//mediapipe/framework/deps:clock",
-        "//mediapipe/framework/port:advanced_proto_lite",
-        "//mediapipe/framework/port:integral_types",
+        "//mediapipe/framework/tool:tag_map",
+        "//mediapipe/framework/tool:validate_name",
         "//mediapipe/framework/port:logging",
         "//mediapipe/framework/port:re2",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:status",
+        "//mediapipe/framework/port:advanced_proto_lite",
         "//mediapipe/framework/tool:name_util",
-        "//mediapipe/framework/tool:tag_map",
-        "//mediapipe/framework/tool:validate_name",
-        "@com_google_absl//absl/memory",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/synchronization",
-        "@com_google_absl//absl/time",
-        "@com_google_absl//absl/types:optional",
-    ],
+    ] + select({
+        "//conditions:default": [],
+    }) + select({
+        "//conditions:default": [
+        ],
+        "//mediapipe/gpu:disable_gpu": [],
+    }),
 )
 
 cc_library(
@@ -270,6 +278,7 @@ cc_library(
     name = "profiler_resource_util",
     srcs = ["profiler_resource_util_common.cc"] + select({
         "//conditions:default": ["profiler_resource_util.cc"],
+        "//mediapipe/framework:android_no_jni": ["profiler_resource_util_android_hal.cc"],
         "//mediapipe:android": ["profiler_resource_util_android.cc"],
         "//mediapipe:ios": ["profiler_resource_util_ios.cc"],
     }),
@@ -295,6 +304,9 @@ cc_library(
         "//conditions:default": [
             "//mediapipe/framework/port:file_helpers",
         ],
+        "//mediapipe/framework:android_no_jni": [
+            "//mediapipe/framework/port:file_helpers",
+        ],
         "//mediapipe:android": [
             "//mediapipe/java/com/google/mediapipe/framework/jni:jni_util",
             "//mediapipe/framework/port:file_helpers",

mediapipe/framework/profiler/graph_profiler.cc

@@ -600,12 +600,17 @@ void AssignNodeNames(GraphProfile* profile) {
   if (graph_trace) {
     graph_trace->clear_calculator_name();
   }
+  std::vector<std::string> canonical_names;
+  canonical_names.reserve(graph_config->node().size());
   for (int i = 0; i < graph_config->node().size(); ++i) {
-    std::string node_name = CanonicalNodeName(*graph_config, i);
-    graph_config->mutable_node(i)->set_name(node_name);
-    if (graph_trace) {
-      graph_trace->add_calculator_name(node_name);
-    }
+    canonical_names.push_back(CanonicalNodeName(*graph_config, i));
+  }
+  for (int i = 0; i < graph_config->node().size(); ++i) {
+    graph_config->mutable_node(i)->set_name(canonical_names[i]);
+  }
+  if (graph_trace) {
+    graph_trace->mutable_calculator_name()->Assign(canonical_names.begin(),
+                                                   canonical_names.end());
   }
 }
 
@@ -646,7 +651,8 @@ absl::StatusOr<std::string> GraphProfiler::GetTraceLogPath() {
   }
 }
 
-absl::Status GraphProfiler::CaptureProfile(GraphProfile* result) {
+absl::Status GraphProfiler::CaptureProfile(
+    GraphProfile* result, PopulateGraphConfig populate_config) {
   // Record the GraphTrace events since the previous WriteProfile.
   // The end_time is chosen to be trace_log_margin_usec in the past,
   // providing time for events to be appended to the TraceBuffer.
@@ -674,6 +680,10 @@ absl::Status GraphProfiler::CaptureProfile(GraphProfile* result) {
   }
   this->Reset();
   CleanCalculatorProfiles(result);
+  if (populate_config == PopulateGraphConfig::kFull) {
+    *result->mutable_config() = validated_graph_->Config();
+    AssignNodeNames(result);
+  }
   return status;
 }
 
@@ -686,7 +696,7 @@ absl::Status GraphProfiler::WriteProfile() {
   int log_interval_count = GetLogIntervalCount(profiler_config_);
   int log_file_count = GetLogFileCount(profiler_config_);
   GraphProfile profile;
-  MP_RETURN_IF_ERROR(CaptureProfile(&profile));
+  MP_RETURN_IF_ERROR(CaptureProfile(&profile, PopulateGraphConfig::kNo));
 
   // If there are no trace events, skip log writing.
   const GraphTrace& trace = *profile.graph_trace().rbegin();

mediapipe/framework/profiler/graph_profiler.h

@@ -71,6 +71,9 @@ struct PacketInfo {
 // For testing
 class GraphProfilerTestPeer;
 
+// GraphProfiler::CaptureProfile option, see the method for details.
+enum class PopulateGraphConfig { kNo, kFull };
+
 // GraphProfiler keeps track of the following in microseconds based on the
 // profiler clock, for each calculator
 // - Open(), Process(), and Close() runtime.
@@ -145,7 +148,14 @@ class GraphProfiler : public std::enable_shared_from_this<ProfilingContext> {
 
   // Records recent profiling and tracing data.  Includes events since the
   // previous call to CaptureProfile.
-  absl::Status CaptureProfile(GraphProfile* result);
+  //
+  // If `populate_config` is `kFull`, `config` field of the resulting profile
+  // will contain canonicalized config of the profiled graph, and
+  // `graph_trace.calculator_name` will contain node names referring to that
+  // config. Both fields are left empty if the option is set to `kNo`.
+  absl::Status CaptureProfile(
+      GraphProfile* result,
+      PopulateGraphConfig populate_config = PopulateGraphConfig::kNo);
 
   // Writes recent profiling and tracing data to a file specified in the
   // ProfilerConfig.  Includes events since the previous call to WriteProfile.
@@ -356,6 +366,85 @@ class ProfilingContext : public GraphProfiler {
 
 // For now, OSS always uses GlContextProfilerStub.
 // TODO: Switch to GlContextProfiler when GlContext is moved to OSS.
+#define MEDIAPIPE_DISABLE_GPU_PROFILER 1
+
+// GlContextProfiler keeps track of all timestamp queries within a specific
+// GlContext object. When created, the GlContextProfiler must be initialized
+// before marking timestamps. Finally, when GlContext is no longer interested
+// in marking timestamps or is about to be destroyed, Finish() must be called
+// to complete all pending time queries and detach the timer from the GlContext.
+// Note that the GlContextProfiler must be created and initialized within a
+// valid GlContext object.
+#if !MEDIAPIPE_DISABLE_GPU_PROFILER
+class GlContextProfiler {
+ public:
+  explicit GlContextProfiler(
+      std::shared_ptr<ProfilingContext> profiling_context)
+      : profiling_context_(profiling_context) {}
+
+  // Not copyable or movable.
+  GlContextProfiler(const GlContextProfiler&) = delete;
+  GlContextProfiler& operator=(const GlContextProfiler&) = delete;
+
+  // Add a GlTimingInfo object to the collection of pending timestamp queries
+  // associated with a specific graph node_id, packet input_timestamp and mark
+  // if it is a start or stop event. When a stop event is marked, this function
+  // blocks on the corresponding start event to complete.
+  void MarkTimestamp(int node_id, Timestamp input_timestamp, bool is_finish);
+
+  // Complete all pending timing queries and detach the timer from the
+  // GlContext.
+  void LogAllTimestamps();
+
+ private:
+  // Store GlTimeQuery and the corresponding TraceEvent object that should be
+  // populated when the query completes together.
+  struct GlTimingInfo {
+    GlTimeQuery time_query;
+    TraceEvent trace_event;
+  };
+
+  // Setup the timer for marking GPU timestamps. If successful in setup, return
+  // true otherwise return false to indicate that timing measurment is not
+  // supported.
+  bool Initialize();
+
+  absl::Time TimeNow();
+
+  // Calibrate the GPU timer w.r.t. the CPU clock. If calibration is fails,
+  // timing_measurement_supported_ is set to false.
+  void CalibrateTimer(bool recalibrate);
+
+  // Log a TraceEvent object to represent if the GPU calibration period has
+  // started or just ended.
+  void LogCalibrationEvent(bool started, absl::Time time);
+
+  // Log TraceEvent objects for completed time queries. If the parameter wait is
+  // set to true, wait for all time queries to complete before returning.
+  void RetireReadyGlTimings(bool wait = false);
+
+  // Get the TraceEvent object containing the timestamp recorded by the GPU if
+  // the provided query was fulfilled. If it is still pending and wait is false,
+  // return absl::nullopt.
+  absl::optional<TraceEvent> GetTimeFromQuery(
+      std::unique_ptr<GlTimingInfo>& query, bool wait);
+
+  std::shared_ptr<ProfilingContext> profiling_context_;
+  GlSimpleTimer gl_timer_;
+  bool checked_timing_supported_ = false;
+  bool timing_measurement_supported_ = false;
+  std::deque<std::unique_ptr<GlTimingInfo>> pending_gl_times_;
+  std::unique_ptr<GlTimingInfo> gl_start_query_;
+};
+
+// The API class used to access the preferred GlContext profiler, such as
+// GlContextProfiler or GlContextProfilerStub. GlProfilingHelper is defined as
+// a class rather than a typedef in order to support clients that refer
+// to it only as a forward declaration.
+class GlProfilingHelper : public GlContextProfiler {
+  using GlContextProfiler::GlContextProfiler;
+};
+#else   // MEDIAPIPE_DISABLE_GPU_PROFILER
 class GlContextProfilerStub {
  public:
   explicit GlContextProfilerStub(
@@ -370,7 +459,8 @@ class GlContextProfilerStub {
 class GlProfilingHelper : public GlContextProfilerStub {
   using GlContextProfilerStub::GlContextProfilerStub;
 };
-
+#endif  // !MEDIAPIPE_DISABLE_GPU_PROFILER
+#undef MEDIAPIPE_DISABLE_GPU_PROFILER
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_FRAMEWORK_PROFILER_GRAPH_PROFILER_H_

mediapipe/framework/profiler/graph_profiler_stub.h

@@ -74,6 +74,9 @@ class TraceEvent {
   inline TraceEvent& set_event_data(int64 data) { return *this; }
 };
 
+// GraphProfiler::CaptureProfile option, see the method for details.
+enum class PopulateGraphConfig { kNo, kFull };
+
 // Empty implementation of ProfilingContext to be used in place of the
 // GraphProfiler when the main implementation is disabled.
 class GraphProfilerStub {
@@ -85,6 +88,11 @@ class GraphProfilerStub {
       std::vector<CalculatorProfile>*) const {
     return absl::OkStatus();
   }
+  absl::Status CaptureProfile(
+      GraphProfile* result,
+      PopulateGraphConfig populate_config = PopulateGraphConfig::kNo) {
+    return absl::OkStatus();
+  }
   inline void Pause() {}
   inline void Resume() {}
   inline void Reset() {}

mediapipe/framework/profiler/graph_profiler_test.cc

@@ -1267,5 +1267,49 @@ TEST(GraphProfilerTest, CalculatorProfileFilter) {
   EXPECT_EQ(GetCalculatorNames(config), expected_names);
 }
 
+TEST(GraphProfilerTest, CaptureProfilePopulateConfig) {
+  CalculatorGraphConfig config;
+  QCHECK(proto2::TextFormat::ParseFromString(R"(
+    profiler_config {
+      enable_profiler: true
+      trace_enabled: true
+    }
+    input_stream: "input_stream"
+    node {
+      calculator: "DummyTestCalculator"
+      input_stream: "input_stream"
+    }
+    node {
+      calculator: "DummyTestCalculator"
+      input_stream: "input_stream"
+    }
+    )",
+                                             &config));
+  CalculatorGraph graph;
+  MP_ASSERT_OK(graph.Initialize(config));
+  GraphProfile profile;
+  MP_ASSERT_OK(
+      graph.profiler()->CaptureProfile(&profile, PopulateGraphConfig::kFull));
+  EXPECT_THAT(profile.config(), Partially(EqualsProto(R"pb(
+                input_stream: "input_stream"
+                node {
+                  name: "DummyTestCalculator_1"
+                  calculator: "DummyTestCalculator"
+                  input_stream: "input_stream"
+                }
+                node {
+                  name: "DummyTestCalculator_2"
+                  calculator: "DummyTestCalculator"
+                  input_stream: "input_stream"
+                }
+              )pb")));
+  EXPECT_THAT(profile.graph_trace(),
+              ElementsAre(Partially(EqualsProto(
+                  R"pb(
+                    calculator_name: "DummyTestCalculator_1"
+                    calculator_name: "DummyTestCalculator_2"
+                  )pb"))));
+}
+
 }  // namespace
 }  // namespace mediapipe

mediapipe/framework/profiler/graph_tracer_test.cc

@@ -1053,8 +1053,8 @@ TEST_F(GraphTracerE2ETest, DemuxGraphLogFiles) {
                   calculator_name: "LambdaCalculator_1"
                   calculator_name: "FlowLimiterCalculator"
                   calculator_name: "RoundRobinDemuxCalculator"
-                  calculator_name: "LambdaCalculator_1"
-                  calculator_name: "LambdaCalculator"
+                  calculator_name: "LambdaCalculator_2"
+                  calculator_name: "LambdaCalculator_3"
                   calculator_name: "ImmediateMuxCalculator"
                   stream_name: ""
                   stream_name: "input_packets_0"
@@ -1198,14 +1198,14 @@ TEST_F(GraphTracerE2ETest, DemuxGraphLogFiles) {
                     output_stream: "OUTPUT:1:input_1"
                   }
                   node {
-                    name: "LambdaCalculator_1"
+                    name: "LambdaCalculator_2"
                     calculator: "LambdaCalculator"
                     input_stream: "input_0"
                     output_stream: "output_0"
                     input_side_packet: "callback_0"
                   }
                   node {
-                    name: "LambdaCalculator"
+                    name: "LambdaCalculator_3"
                     calculator: "LambdaCalculator"
                     input_stream: "input_1"
                     output_stream: "output_1"

mediapipe/framework/profiler/profiler_resource_util_android_hal.cc

@@ -0,0 +1,9 @@
+#include "mediapipe/framework/port/statusor.h"
+
+namespace mediapipe {
+
+StatusOr<std::string> GetDefaultTraceLogDirectory() {
+  return "/data/local/tmp";
+}
+
+}  // namespace mediapipe

mediapipe/framework/profiler/reporter/print_profile.cc

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
-// This program takes one input file and encodes its contents as a C++
-// std::string, which can be included in a C++ source file. It is similar to
-// filewrapper (and borrows some of its code), but simpler.
+// This program takes one input file and encodes its contents as a C++ string,
+// which can be included in a C++ source file. It is similar to filewrapper
+// (and borrows some of its code), but simpler.
 
 #include <algorithm>
 #include <fstream>

mediapipe/framework/profiler/reporter/reporter.cc

@@ -215,7 +215,7 @@ void CompleteCalculatorData(
 }
 
 void Reporter::Accumulate(const mediapipe::GraphProfile& profile) {
-  // Cache nodeID to its std::string name.
+  // Cache nodeID to its string name.
   NameLookup name_lookup;
   CacheNodeNameLookup(profile, &name_lookup);
 
@@ -363,8 +363,8 @@ class ReportImpl : public Report {
   // Values for each calculator, corresponding to the label in headers().
   std::vector<std::vector<std::string>> lines_impl;
 
-  // The longest std::string of any value in a given column (including the
-  // header for that column). Used for formatting the output.
+  // The longest string of any value in a given column (including the header
+  // for that column). Used for formatting the output.
   std::vector<size_t> char_counts_impl;
   bool compact_flag = false;
 
@@ -377,7 +377,7 @@ void ReportImpl::Print(std::ostream& output) {
   // fill space up to char_counts[column] + 1. The strings in the output
   // are mutable to support padding, hence no const in the for loops.
   int column_number = 0;
-  // Make a copy of the column std::string because we might be adding spaces.
+  // Make a copy of the column string because we might be adding spaces.
   for (auto column : headers_impl) {
     int padding_needed = char_counts_impl[column_number] + 1 - column.length();
     if (compact_flag) {

mediapipe/framework/profiler/trace_builder.cc

@@ -90,11 +90,11 @@ void BasicTraceEventTypes(TraceEventRegistry* result) {
   }
 }
 
-// A map defining int32 identifiers for std::string object pointers.
-// Lookup is fast when the same std::string object is used frequently.
+// A map defining int32 identifiers for string object pointers.
+// Lookup is fast when the same string object is used frequently.
 class StringIdMap {
  public:
-  // Returns the int32 identifier for a std::string object pointer.
+  // Returns the int32 identifier for a string object pointer.
   int32 operator[](const std::string* id) {
     if (id == nullptr) {
       return 0;

mediapipe/framework/stream_handler/sync_set_input_stream_handler_test.cc

@@ -47,8 +47,8 @@ std::tuple<std::string, Timestamp, std::vector<std::string>> CommandTuple(
   return std::make_tuple(stream, timestamp, expected);
 }
 
-// Function to take the inputs and produce a diagnostic output std::string
-// and output a packet with a diagnostic output std::string which includes
+// Function to take the inputs and produce a diagnostic output string
+// and output a packet with a diagnostic output string which includes
 // the input timestamp and the ids of each input which is present.
 absl::Status InputsToDebugString(const InputStreamShardSet& inputs,
                                  OutputStreamShardSet* outputs) {

mediapipe/framework/test_calculators.cc

@@ -279,7 +279,7 @@ class StdDevCalculator : public CalculatorBase {
 REGISTER_CALCULATOR(StdDevCalculator);
 
 // A calculator that receives some number of input streams carrying ints.
-// Outputs, for each input timestamp, a space separated std::string containing
+// Outputs, for each input timestamp, a space separated string containing
 // the timestamp and all the inputs for that timestamp (Empty inputs
 // will be denoted with "empty"). Sets the header to be a space-separated
 // concatenation of the input stream headers.
@@ -368,7 +368,7 @@ REGISTER_CALCULATOR(SaverCalculator);
 
 #ifndef MEDIAPIPE_MOBILE
 // Source Calculator that produces matrices on the output stream with
-// each coefficient from a normal gaussian.  A std::string seed must be given
+// each coefficient from a normal gaussian.  A string seed must be given
 // as an input side packet.
 class RandomMatrixCalculator : public CalculatorBase {
  public:

mediapipe/framework/timestamp.h

@@ -86,7 +86,7 @@ class Timestamp {
   // in microseconds, but this function should be preferred over Value() in case
   // the underlying representation changes.
   int64 Microseconds() const { return Value(); }
-  // This provides a human readable std::string for the special values.
+  // This provides a human readable string for the special values.
   std::string DebugString() const;
 
   // For use by framework. Clients or Calculator implementations should not call

mediapipe/framework/tool/BUILD

@@ -43,6 +43,7 @@ bzl_library(
         "//mediapipe/framework:transitive_protos_bzl",
         "//mediapipe/framework/deps:descriptor_set_bzl",
         "//mediapipe/framework/deps:expand_template_bzl",
+        "@org_tensorflow//tensorflow/lite/core/shims:cc_library_with_tflite_bzl",
     ],
 )
 
@@ -52,6 +53,7 @@ bzl_library(
         "build_defs.bzl",
     ],
     visibility = [
+        "//mediapipe/app/xeno/catalog:__subpackages__",
         "//mediapipe/framework:__subpackages__",
     ],
 )
@@ -286,6 +288,7 @@ cc_library(
 mediapipe_cc_test(
     name = "options_util_test",
     size = "small",
+    timeout = "moderate",
     srcs = ["options_util_test.cc"],
     # A non-empty "data" param is needed to build the "_test_wasm" target.
     data = [":node_chain_subgraph.proto"],
@@ -858,11 +861,9 @@ cc_library(
         "//mediapipe/framework:calculator_cc_proto",
         "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework:mediapipe_options_cc_proto",
-        "//mediapipe/framework:stream_handler_cc_proto",
         "//mediapipe/framework:subgraph",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:status",
-        "//mediapipe/framework/stream_handler:sync_set_input_stream_handler_cc_proto",
         "//mediapipe/framework/tool:switch_container_cc_proto",
     ],
     alwayslink = 1,

mediapipe/framework/tool/encode_as_c_string.cc

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
-// This program takes one input file and encodes its contents as a C++
-// std::string, which can be included in a C++ source file. It is similar to
-// filewrapper (and borrows some of its code), but simpler.
+// This program takes one input file and encodes its contents as a C++ string,
+// which can be included in a C++ source file. It is similar to filewrapper
+// (and borrows some of its code), but simpler.
 
 #include <fstream>
 #include <iostream>

mediapipe/framework/tool/mediapipe_graph.bzl

@@ -20,6 +20,7 @@ load("//mediapipe/framework:transitive_protos.bzl", "transitive_protos")
 load("//mediapipe/framework/deps:expand_template.bzl", "expand_template")
 load("//mediapipe/framework/tool:build_defs.bzl", "clean_dep")
 load("//mediapipe/framework/deps:descriptor_set.bzl", "direct_descriptor_set", "transitive_descriptor_set")
+load("@org_tensorflow//tensorflow/lite/core/shims:cc_library_with_tflite.bzl", "cc_library_with_tflite")
 
 def mediapipe_binary_graph(name, graph = None, output_name = None, deps = [], testonly = False, **kwargs):
     """Converts a graph from text format to binary format."""
@@ -98,6 +99,7 @@ def mediapipe_simple_subgraph(
         register_as,
         graph,
         deps = [],
+        tflite_deps = None,
         visibility = None,
         testonly = None,
         **kwargs):
@@ -109,6 +111,7 @@ def mediapipe_simple_subgraph(
           CamelCase.
       graph: the BUILD label of a text-format MediaPipe graph.
       deps: any calculators or subgraphs used by this graph.
+      tflite_deps: any calculators or subgraphs used by this graph that may use different TFLite implementation.
       visibility: The list of packages the subgraph should be visible to.
       testonly: pass 1 if the graph is to be used only for tests.
       **kwargs: Remaining keyword args, forwarded to cc_library.
@@ -138,21 +141,39 @@ def mediapipe_simple_subgraph(
         },
         testonly = testonly,
     )
-    native.cc_library(
-        name = name,
-        srcs = [
-            name + "_linked.cc",
-            graph_base_name + ".inc",
-        ],
-        deps = [
-            clean_dep("//mediapipe/framework:calculator_framework"),
-            clean_dep("//mediapipe/framework:subgraph"),
-        ] + deps,
-        alwayslink = 1,
-        visibility = visibility,
-        testonly = testonly,
-        **kwargs
-    )
+    if not tflite_deps:
+        native.cc_library(
+            name = name,
+            srcs = [
+                name + "_linked.cc",
+                graph_base_name + ".inc",
+            ],
+            deps = [
+                clean_dep("//mediapipe/framework:calculator_framework"),
+                clean_dep("//mediapipe/framework:subgraph"),
+            ] + deps,
+            alwayslink = 1,
+            visibility = visibility,
+            testonly = testonly,
+            **kwargs
+        )
+    else:
+        cc_library_with_tflite(
+            name = name,
+            srcs = [
+                name + "_linked.cc",
+                graph_base_name + ".inc",
+            ],
+            tflite_deps = tflite_deps,
+            deps = [
+                clean_dep("//mediapipe/framework:calculator_framework"),
+                clean_dep("//mediapipe/framework:subgraph"),
+            ] + deps,
+            alwayslink = 1,
+            visibility = visibility,
+            testonly = testonly,
+            **kwargs
+        )
 
 def mediapipe_reexport_library(
         name,

mediapipe/framework/tool/name_util.h

@@ -85,7 +85,7 @@ std::pair<std::string, int> ParseTagIndexFromStream(const std::string& stream);
 // Formats to "tag:index".
 std::string CatTag(const std::string& tag, int index);
 
-// Concatenates "tag:index:name" into a single std::string.
+// Concatenates "tag:index:name" into a single string.
 std::string CatStream(const std::pair<std::string, int>& tag_index,
                       const std::string& name);
 

mediapipe/framework/tool/options_syntax_util.h

@@ -34,7 +34,7 @@ class OptionsSyntaxUtil {
   FieldPath OptionFieldPath(absl::string_view tag,
                             const Descriptor* descriptor);
 
-  // Splits a std::string into "tag" and "name" delimited by a single colon.
+  // Splits a string into "tag" and "name" delimited by a single colon.
   std::vector<absl::string_view> StrSplitTags(absl::string_view tag_and_name);
 
  private:

mediapipe/framework/tool/proto_util_lite.cc

@@ -228,7 +228,7 @@ absl::Status SyntaxStatus(bool ok, const std::string& text, T* result) {
                   " for type: ", MediaPipeTypeStringOrDemangled<T>(), "."));
 }
 
-// Templated parsing of a std::string value.
+// Templated parsing of a string value.
 template <typename T>
 absl::Status ParseValue(const std::string& text, T* result) {
   return SyntaxStatus(absl::SimpleAtoi(text, result), text, result);

mediapipe/framework/tool/simple_subgraph_template.cc

@@ -31,7 +31,7 @@ class {{SUBGRAPH_CLASS_NAME}} : public Subgraph {
         const SubgraphOptions& /*options*/) {
     CalculatorGraphConfig config;
     // Note: this is a binary protobuf serialization, and may include NUL
-    // bytes. The trailing NUL added to the std::string literal should be excluded.
+    // bytes. The trailing NUL added to the string literal should be excluded.
     if (config.ParseFromArray(binary_graph, sizeof(binary_graph) - 1)) {
       return config;
     } else {

mediapipe/framework/tool/status_util.h

@@ -40,7 +40,7 @@ absl::Status StatusInvalid(const std::string& error_message);
 ABSL_DEPRECATED("Use absl::UnknownError(error_message) instead.")
 absl::Status StatusFail(const std::string& error_message);
 
-// Prefixes the given std::string to the error message in status.
+// Prefixes the given string to the error message in status.
 // This function should be considered internal to the framework.
 // TODO Replace usage of AddStatusPrefix with util::Annotate().
 absl::Status AddStatusPrefix(const std::string& prefix,

mediapipe/framework/tool/switch_container.cc

@@ -22,8 +22,6 @@
 #include "mediapipe/framework/port/canonical_errors.h"
 #include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/framework/port/status.h"
-#include "mediapipe/framework/stream_handler.pb.h"
-#include "mediapipe/framework/stream_handler/sync_set_input_stream_handler.pb.h"
 #include "mediapipe/framework/tool/container_util.h"
 #include "mediapipe/framework/tool/name_util.h"
 #include "mediapipe/framework/tool/subgraph_expansion.h"
@@ -88,12 +86,6 @@ CalculatorGraphConfig::Node* BuildDemuxNode(
     CalculatorGraphConfig* config) {
   CalculatorGraphConfig::Node* result = config->add_node();
   *result->mutable_calculator() = "SwitchDemuxCalculator";
-  *result->mutable_input_stream_handler()->mutable_input_stream_handler() =
-      "ImmediateInputStreamHandler";
-  if (container_node.has_input_stream_handler()) {
-    *result->mutable_input_stream_handler() =
-        container_node.input_stream_handler();
-  }
   return result;
 }
 
@@ -103,8 +95,6 @@ CalculatorGraphConfig::Node* BuildMuxNode(
     CalculatorGraphConfig* config) {
   CalculatorGraphConfig::Node* result = config->add_node();
   *result->mutable_calculator() = "SwitchMuxCalculator";
-  *result->mutable_input_stream_handler()->mutable_input_stream_handler() =
-      "ImmediateInputStreamHandler";
   return result;
 }
 

mediapipe/framework/tool/switch_container.proto

@@ -24,4 +24,7 @@ message SwitchContainerOptions {
 
   // Activates channel 1 for enable = true, channel 0 otherwise.
   optional bool enable = 4;
+
+  // Use DefaultInputStreamHandler for muxing & demuxing.
+  optional bool synchronize_io = 5;
 }

mediapipe/framework/tool/switch_container_test.cc

@@ -66,8 +66,7 @@ REGISTER_CALCULATOR(TripleIntCalculator);
 // A testing example of a SwitchContainer containing two subnodes.
 // Note that the input and output tags supplied to the container node,
 // must match the input and output tags required by the subnodes.
-CalculatorGraphConfig SubnodeContainerExample(
-    const std::string& input_stream_handler = "") {
+CalculatorGraphConfig SubnodeContainerExample(const std::string& options = "") {
   std::string config = R"pb(
     input_stream: "foo"
     input_stream: "enable"
@@ -80,9 +79,9 @@ CalculatorGraphConfig SubnodeContainerExample(
       options {
         [mediapipe.SwitchContainerOptions.ext] {
           contained_node: { calculator: "TripleIntCalculator" }
-          contained_node: { calculator: "PassThroughCalculator" }
+          contained_node: { calculator: "PassThroughCalculator" } $options
         }
-      } $input_stream_handler
+      }
     }
     node {
       calculator: "PassThroughCalculator"
@@ -94,8 +93,7 @@ CalculatorGraphConfig SubnodeContainerExample(
   )pb";
 
   return mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
-      absl::StrReplaceAll(config,
-                          {{"$input_stream_handler", input_stream_handler}}));
+      absl::StrReplaceAll(config, {{"$options", options}}));
 }
 
 // A testing example of a SwitchContainer containing two subnodes.
@@ -248,9 +246,6 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
           options {
             [mediapipe.SwitchContainerOptions.ext] {}
           }
-          input_stream_handler {
-            input_stream_handler: "ImmediateInputStreamHandler"
-          }
         }
         node {
           name: "switchcontainer__TripleIntCalculator"
@@ -274,9 +269,6 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
           options {
             [mediapipe.SwitchContainerOptions.ext] {}
           }
-          input_stream_handler {
-            input_stream_handler: "ImmediateInputStreamHandler"
-          }
         }
         node {
           calculator: "PassThroughCalculator"
@@ -322,9 +314,7 @@ TEST(SwitchContainerTest, ValidateInputStreamHandler) {
       options {
         [mediapipe.SwitchContainerOptions.ext] {}
       }
-      input_stream_handler {
-        input_stream_handler: "ImmediateInputStreamHandler"
-      }
+      input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
     }
     node {
       name: "switchcontainer__TripleIntCalculator"
@@ -350,9 +340,7 @@ TEST(SwitchContainerTest, ValidateInputStreamHandler) {
       options {
         [mediapipe.SwitchContainerOptions.ext] {}
       }
-      input_stream_handler {
-        input_stream_handler: "ImmediateInputStreamHandler"
-      }
+      input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
     }
     node {
       calculator: "PassThroughCalculator"
@@ -371,83 +359,12 @@ TEST(SwitchContainerTest, ValidateInputStreamHandler) {
   EXPECT_THAT(graph.Config(), mediapipe::EqualsProto(expected_graph));
 }
 
-// Expands the SwitchContainer with a node-level input_stream_handler.
-TEST(SwitchContainerTest, OverrideInputStreamHandler) {
-  EXPECT_TRUE(SubgraphRegistry::IsRegistered("SwitchContainer"));
-  CalculatorGraph graph;
-  CalculatorGraphConfig supergraph = SubnodeContainerExample(
-      R"pb(input_stream_handler {
-             input_stream_handler: "DefaultInputStreamHandler"
-           })pb");
-  *supergraph.mutable_node(0)
-       ->mutable_input_stream_handler()
-       ->mutable_input_stream_handler() = "DefaultInputStreamHandler";
-  MP_ASSERT_OK(graph.Initialize(supergraph, {}));
-  CalculatorGraphConfig expected_graph =
-      mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
-        node {
-          name: "switchcontainer__SwitchDemuxCalculator"
-          calculator: "SwitchDemuxCalculator"
-          input_stream: "ENABLE:enable"
-          input_stream: "foo"
-          output_stream: "C0__:switchcontainer__c0__foo"
-          output_stream: "C1__:switchcontainer__c1__foo"
-          options {
-            [mediapipe.SwitchContainerOptions.ext] {}
-          }
-          input_stream_handler {
-            input_stream_handler: "DefaultInputStreamHandler"
-          }
-        }
-        node {
-          name: "switchcontainer__TripleIntCalculator"
-          calculator: "TripleIntCalculator"
-          input_stream: "switchcontainer__c0__foo"
-          output_stream: "switchcontainer__c0__bar"
-        }
-        node {
-          name: "switchcontainer__PassThroughCalculator"
-          calculator: "PassThroughCalculator"
-          input_stream: "switchcontainer__c1__foo"
-          output_stream: "switchcontainer__c1__bar"
-        }
-        node {
-          name: "switchcontainer__SwitchMuxCalculator"
-          calculator: "SwitchMuxCalculator"
-          input_stream: "ENABLE:enable"
-          input_stream: "C0__:switchcontainer__c0__bar"
-          input_stream: "C1__:switchcontainer__c1__bar"
-          output_stream: "bar"
-          options {
-            [mediapipe.SwitchContainerOptions.ext] {}
-          }
-          input_stream_handler {
-            input_stream_handler: "ImmediateInputStreamHandler"
-          }
-        }
-        node {
-          calculator: "PassThroughCalculator"
-          input_stream: "foo"
-          input_stream: "bar"
-          output_stream: "output_foo"
-          output_stream: "output_bar"
-        }
-        input_stream: "foo"
-        input_stream: "enable"
-        executor {}
-        input_side_packet: "timezone"
-      )pb");
-  EXPECT_THAT(graph.Config(), mediapipe::EqualsProto(expected_graph));
-}
-
-// Runs the SwitchContainer with a node-level input_stream_handler.
 TEST(SwitchContainerTest, RunsWithInputStreamHandler) {
   EXPECT_TRUE(SubgraphRegistry::IsRegistered("SwitchContainer"));
-  CalculatorGraphConfig supergraph = SubnodeContainerExample(
-      R"pb(input_stream_handler {
-             input_stream_handler: "DefaultInputStreamHandler"
-           })pb");
+  CalculatorGraphConfig supergraph =
+      SubnodeContainerExample(R"pb(synchronize_io: true)pb");
   MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
+  LOG(INFO) << supergraph.DebugString();
   RunTestContainer(supergraph, true);
 }
 
@@ -470,9 +387,6 @@ TEST(SwitchContainerTest, ApplyToSideSubnodes) {
           options {
             [mediapipe.SwitchContainerOptions.ext] {}
           }
-          input_stream_handler {
-            input_stream_handler: "ImmediateInputStreamHandler"
-          }
         }
         node {
           name: "switchcontainer__TripleIntCalculator"
@@ -496,9 +410,6 @@ TEST(SwitchContainerTest, ApplyToSideSubnodes) {
           options {
             [mediapipe.SwitchContainerOptions.ext] {}
           }
-          input_stream_handler {
-            input_stream_handler: "ImmediateInputStreamHandler"
-          }
         }
         node {
           calculator: "PassThroughCalculator"

mediapipe/framework/tool/switch_demux_calculator.cc

@@ -26,6 +26,7 @@
 #include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/port/status_macros.h"
 #include "mediapipe/framework/tool/container_util.h"
+#include "mediapipe/framework/tool/switch_container.pb.h"
 
 namespace mediapipe {
 
@@ -113,7 +114,10 @@ absl::Status SwitchDemuxCalculator::GetContract(CalculatorContract* cc) {
       }
     }
   }
-  cc->SetInputStreamHandler("ImmediateInputStreamHandler");
+  auto& options = cc->Options<mediapipe::SwitchContainerOptions>();
+  if (!options.synchronize_io()) {
+    cc->SetInputStreamHandler("ImmediateInputStreamHandler");
+  }
   cc->SetProcessTimestampBounds(true);
   return absl::OkStatus();
 }

mediapipe/framework/tool/switch_mux_calculator.cc

@@ -28,6 +28,7 @@
 #include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/port/status_macros.h"
 #include "mediapipe/framework/tool/container_util.h"
+#include "mediapipe/framework/tool/switch_container.pb.h"
 
 namespace mediapipe {
 
@@ -68,6 +69,17 @@ class SwitchMuxCalculator : public CalculatorBase {
  private:
   int channel_index_;
   std::set<std::string> channel_tags_;
+  mediapipe::SwitchContainerOptions options_;
+  // This is used to keep around packets that we've received but not
+  // relayed yet (because we may not know which channel we should yet be using
+  // when synchronized_io flag is set).
+  std::map<Timestamp, std::map<CollectionItemId, Packet>> packet_history_;
+  // Historical channel index values for timestamps where we don't have all
+  // packets available yet (when synchronized_io flag is set).
+  std::map<Timestamp, int> channel_history_;
+  // Number of output steams that we already processed for the current output
+  // timestamp.
+  int current_processed_stream_count_ = 0;
 };
 REGISTER_CALCULATOR(SwitchMuxCalculator);
 
@@ -122,6 +134,7 @@ absl::Status SwitchMuxCalculator::GetContract(CalculatorContract* cc) {
 }
 
 absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) {
+  options_ = cc->Options<mediapipe::SwitchContainerOptions>();
   channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
   channel_tags_ = ChannelTags(cc->Inputs().TagMap());
 
@@ -141,13 +154,79 @@ absl::Status SwitchMuxCalculator::Process(CalculatorContext* cc) {
   // Update the input channel index if specified.
   channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
 
-  // Relay packets and timestamps only from channel_index_.
-  for (const std::string& tag : channel_tags_) {
-    for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
-      auto& output = cc->Outputs().Get(tag, index);
-      std::string input_tag = tool::ChannelTag(tag, channel_index_);
-      auto& input = cc->Inputs().Get(input_tag, index);
-      tool::Relay(input, &output);
+  if (options_.synchronize_io()) {
+    // Start with adding input signals into channel_history_ and packet_history_
+    if (cc->Inputs().HasTag("ENABLE") &&
+        !cc->Inputs().Tag("ENABLE").IsEmpty()) {
+      channel_history_[cc->Inputs().Tag("ENABLE").Value().Timestamp()] =
+          channel_index_;
+    }
+    if (cc->Inputs().HasTag("SELECT") &&
+        !cc->Inputs().Tag("SELECT").IsEmpty()) {
+      channel_history_[cc->Inputs().Tag("SELECT").Value().Timestamp()] =
+          channel_index_;
+    }
+    for (auto input_id = cc->Inputs().BeginId();
+         input_id < cc->Inputs().EndId(); ++input_id) {
+      auto& entry = cc->Inputs().Get(input_id);
+      if (entry.IsEmpty()) {
+        continue;
+      }
+      packet_history_[entry.Value().Timestamp()][input_id] = entry.Value();
+    }
+    // Now check if we have enough information to produce any outputs.
+    while (!channel_history_.empty()) {
+      // Look at the oldest unprocessed timestamp.
+      auto it = channel_history_.begin();
+      auto& packets = packet_history_[it->first];
+      int total_streams = 0;
+      // Loop over all outputs to see if we have anything new that we can relay.
+      for (const std::string& tag : channel_tags_) {
+        for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
+          ++total_streams;
+          auto input_id =
+              cc->Inputs().GetId(tool::ChannelTag(tag, it->second), index);
+          auto packet_it = packets.find(input_id);
+          if (packet_it != packets.end()) {
+            cc->Outputs().Get(tag, index).AddPacket(packet_it->second);
+            ++current_processed_stream_count_;
+          } else if (it->first <
+                     cc->Inputs().Get(input_id).Value().Timestamp()) {
+            // Getting here means that input stream that corresponds to this
+            // output at the timestamp we're trying to process right now has
+            // already advanced beyond this timestamp. This means that we will
+            // shouldn't expect a packet for this timestamp anymore, and we can
+            // safely advance timestamp on the output.
+            cc->Outputs()
+                .Get(tag, index)
+                .SetNextTimestampBound(it->first.NextAllowedInStream());
+            ++current_processed_stream_count_;
+          }
+        }
+      }
+      if (current_processed_stream_count_ == total_streams) {
+        // There's nothing else to wait for at the current timestamp, do the
+        // cleanup and move on to the next one.
+        packet_history_.erase(it->first);
+        channel_history_.erase(it);
+        current_processed_stream_count_ = 0;
+      } else {
+        // We're still missing some packets for the current timestamp. Clean up
+        // those that we just relayed and let the rest wait until the next
+        // Process() call.
+        packets.clear();
+        break;
+      }
+    }
+  } else {
+    // Relay packets and timestamps only from channel_index_.
+    for (const std::string& tag : channel_tags_) {
+      for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
+        auto& output = cc->Outputs().Get(tag, index);
+        std::string input_tag = tool::ChannelTag(tag, channel_index_);
+        auto& input = cc->Inputs().Get(input_tag, index);
+        tool::Relay(input, &output);
+      }
     }
   }
   return absl::OkStatus();

mediapipe/framework/tool/tag_map.h

@@ -51,8 +51,8 @@ class TagMap {
     int count;
   };
 
-  // Create a TagMap from a repeated std::string proto field of
-  // TAG:<index>:name. This is the most common usage:
+  // Create a TagMap from a repeated string proto field of TAG:<index>:name.
+  // This is the most common usage:
   // ASSIGN_OR_RETURN(std::shared_ptr<TagMap> tag_map,
   //                  tool::TagMap::Create(node.input_streams()));
   static absl::StatusOr<std::shared_ptr<TagMap>> Create(
@@ -87,7 +87,7 @@ class TagMap {
 
   // Returns canonicalized strings describing the TagMap.
   proto_ns::RepeatedPtrField<ProtoString> CanonicalEntries() const;
-  // Returns a std::string description for debug purposes.
+  // Returns a string description for debug purposes.
   std::string DebugString() const;
   // Returns a shorter description for debug purposes (doesn't include
   // stream/side packet names).

mediapipe/framework/tool/tag_map_test.cc

@@ -318,8 +318,8 @@ TEST(TagMapTest, SameAs) {
   }
 }
 
-// A helper function to test that a TagMap's debug std::string and short
-// debug std::string each satisfy a matcher.
+// A helper function to test that a TagMap's debug string and short
+// debug string each satisfy a matcher.
 template <typename Matcher>
 void TestDebugString(
     const absl::StatusOr<std::shared_ptr<tool::TagMap>>& statusor_tag_map,

mediapipe/framework/tool/template_expander.cc

@@ -433,7 +433,7 @@ class TemplateExpanderImpl {
     return result;
   }
 
-  // Converts a TemplateArgument to std::string.
+  // Converts a TemplateArgument to string.
   std::string AsString(const TemplateArgument& value) {
     std::string result;
     if (value.has_num()) {

mediapipe/framework/tool/template_parser.cc

@@ -260,7 +260,7 @@ class TemplateParser::Parser::ParserImpl {
   typedef proto_ns::TextFormat::ParseLocation ParseLocation;
 
   // Determines if repeated values for non-repeated fields and
-  // oneofs are permitted, e.g., the std::string "foo: 1 foo: 2" for a
+  // oneofs are permitted, e.g., the string "foo: 1 foo: 2" for a
   // required/optional field named "foo", or "baz: 1 qux: 2"
   // where "baz" and "qux" are members of the same oneof.
   enum SingularOverwritePolicy {
@@ -401,7 +401,7 @@ class TemplateParser::Parser::ParserImpl {
   }
 
 #ifndef PROTO2_OPENSOURCE
-  // Consumes a std::string value and parses it as a packed repeated field into
+  // Consumes a string value and parses it as a packed repeated field into
   // the given field of the given message.
   bool ConsumePackedFieldAsString(const std::string& field_name,
                                   const FieldDescriptor* field,
@@ -409,7 +409,7 @@ class TemplateParser::Parser::ParserImpl {
     std::string packed;
     DO(ConsumeString(&packed));
 
-    // Prepend field tag and varint-encoded std::string length to turn into
+    // Prepend field tag and varint-encoded string length to turn into
     // encoded message.
     std::string tagged;
     {
@@ -428,7 +428,7 @@ class TemplateParser::Parser::ParserImpl {
     io::CodedInputStream coded_input(&array_input);
     if (!message->MergePartialFromCodedStream(&coded_input)) {
       ReportError("Could not parse packed field \"" + field_name +
-                  "\" as wire-encoded std::string.");
+                  "\" as wire-encoded string.");
       return false;
     }
 
@@ -607,7 +607,7 @@ class TemplateParser::Parser::ParserImpl {
       bool consumed_semicolon = TryConsume(":");
       if (consumed_semicolon && field->options().weak() &&
           LookingAtType(io::Tokenizer::TYPE_STRING)) {
-        // we are getting a bytes std::string for a weak field.
+        // we are getting a bytes string for a weak field.
         std::string tmp;
         DO(ConsumeString(&tmp));
         reflection->MutableMessage(message, field)->ParseFromString(tmp);
@@ -640,8 +640,8 @@ class TemplateParser::Parser::ParserImpl {
 #ifndef PROTO2_OPENSOURCE
     } else if (field->is_packable() &&
                LookingAtType(io::Tokenizer::TYPE_STRING)) {
-      // Packable field printed as wire-formatted std::string: "foo: "abc\123"".
-      // Fields of type std::string cannot be packed themselves, so this is
+      // Packable field printed as wire-formatted string: "foo: "abc\123"".
+      // Fields of type string cannot be packed themselves, so this is
       // unambiguous.
       DO(ConsumePackedFieldAsString(field_name, field, message));
 #endif  // !PROTO2_OPENSOURCE
@@ -908,7 +908,7 @@ class TemplateParser::Parser::ParserImpl {
       }
       return true;
     }
-    // Possible field values other than std::string:
+    // Possible field values other than string:
     //   12345        => TYPE_INTEGER
     //   -12345       => TYPE_SYMBOL + TYPE_INTEGER
     //   1.2345       => TYPE_FLOAT
@@ -992,7 +992,7 @@ class TemplateParser::Parser::ParserImpl {
     return false;
   }
 
-  // Consume a std::string of form "<id1>.<id2>....<idN>".
+  // Consume a string of form "<id1>.<id2>....<idN>".
   bool ConsumeFullTypeName(std::string* name) {
     DO(ConsumeIdentifier(name));
     while (TryConsume(".")) {
@@ -1013,11 +1013,11 @@ class TemplateParser::Parser::ParserImpl {
     return true;
   }
 
-  // Consumes a std::string and saves its value in the text parameter.
+  // Consumes a string and saves its value in the text parameter.
   // Returns false if the token is not of type STRING.
   bool ConsumeString(std::string* text) {
     if (!LookingAtType(io::Tokenizer::TYPE_STRING)) {
-      ReportError("Expected std::string, got: " + tokenizer_.current().text);
+      ReportError("Expected string, got: " + tokenizer_.current().text);
       return false;
     }
 
@@ -1391,7 +1391,7 @@ void StowFieldValue(Message* message, TemplateExpression* expression) {
   }
 }
 
-// Strips first and last quotes from a std::string.
+// Strips first and last quotes from a string.
 static void StripQuotes(std::string* str) {
   // Strip off the leading and trailing quotation marks from the value, if
   // there are any.
@@ -1585,7 +1585,7 @@ class TemplateParser::Parser::MediaPipeParserImpl
     return true;
   }
 
-  // Parses a numeric or a std::string literal.
+  // Parses a numeric or a string literal.
   bool ConsumeLiteral(TemplateExpression* result) {
     std::string token = tokenizer_.current().text;
     StripQuotes(&token);

mediapipe/framework/tool/test_util.cc

@@ -215,6 +215,11 @@ std::string GetTestRootDir() {
 std::string GetTestOutputsDir() {
   const char* output_dir = getenv("TEST_UNDECLARED_OUTPUTS_DIR");
   if (!output_dir) {
+#ifdef __APPLE__
+    char path[PATH_MAX];
+    size_t n = confstr(_CS_DARWIN_USER_TEMP_DIR, path, sizeof(path));
+    if (n > 0 && n < sizeof(path)) return path;
+#endif  // __APPLE__
     output_dir = "/tmp";
   }
   return output_dir;

mediapipe/framework/tool/validate_name.h

@@ -66,7 +66,7 @@ absl::Status SetFromTagAndNameInfo(
     const TagAndNameInfo& info,
     proto_ns::RepeatedPtrField<ProtoString>* tags_and_names);
 
-// The std::string is a valid name for an input stream, output stream,
+// The string is a valid name for an input stream, output stream,
 // side packet, and input collection.  Names use only lower case letters,
 // numbers, and underscores.
 //
@@ -77,18 +77,18 @@ absl::Status SetFromTagAndNameInfo(
 //     (3) Because input side packet names end up in model directory names,
 //         where lower case naming is the norm.
 absl::Status ValidateName(const std::string& name);
-// The std::string is a valid tag name.  Tags use only upper case letters,
+// The string is a valid tag name.  Tags use only upper case letters,
 // numbers, and underscores.
 absl::Status ValidateTag(const std::string& tag);
 
-// Parse a "Tag and Name" std::string into a tag and a name.
+// Parse a "Tag and Name" string into a tag and a name.
 // The format is an optional tag and colon, followed by a name.
 // Example 1: "VIDEO:frames2" -> tag: "VIDEO", name: "frames2"
 // Example 2: "video_frames_1" -> tag: "", name: "video_frames_1"
 absl::Status ParseTagAndName(const std::string& tag_and_name, std::string* tag,
                              std::string* name);
 
-// Parse a generic TAG:index:name std::string.  The format is a tag, then an
+// Parse a generic TAG:index:name string.  The format is a tag, then an
 // index, then a name.  The tag and index are optional.  If the index
 // is included, then the tag must be included.  If no tag is used then
 // index is set to -1 (and should be assigned by argument position).
@@ -99,7 +99,7 @@ absl::Status ParseTagAndName(const std::string& tag_and_name, std::string* tag,
 absl::Status ParseTagIndexName(const std::string& tag_and_name,
                                std::string* tag, int* index, std::string* name);
 
-// Parse a generic TAG:index std::string.  The format is a tag, then an index
+// Parse a generic TAG:index string.  The format is a tag, then an index
 // with both being optional.  If the tag is missing it is assumed to be
 // "" and if the index is missing then it is assumed to be 0.  If the
 // index is provided then a colon (':') must be used.

mediapipe/framework/type_map.h

@@ -13,8 +13,8 @@
 // limitations under the License.
 
 // This header defines static maps to store the mappings from type hash id and
-// name std::string to MediaPipeTypeData.  It also provides code to inspect
-// types of packets and access registered serialize and deserialize functions.
+// name string to MediaPipeTypeData.  It also provides code to inspect types of
+// packets and access registered serialize and deserialize functions.
 // Calculators can use this to infer types of packets and adjust accordingly.
 //
 // Register a type:
@@ -242,7 +242,7 @@ class StaticMap {
   class MapName : public type_map_internal::StaticMap<MapName, KeyType> {};
 // Defines a map from unique typeid number to MediaPipeTypeData.
 DEFINE_MEDIAPIPE_TYPE_MAP(PacketTypeIdToMediaPipeTypeData, size_t);
-// Defines a map from unique type std::string to MediaPipeTypeData.
+// Defines a map from unique type string to MediaPipeTypeData.
 DEFINE_MEDIAPIPE_TYPE_MAP(PacketTypeStringToMediaPipeTypeData, std::string);
 
 // MEDIAPIPE_REGISTER_TYPE can be used to register a type.
@@ -267,7 +267,7 @@ DEFINE_MEDIAPIPE_TYPE_MAP(PacketTypeStringToMediaPipeTypeData, std::string);
 //   #undef MY_MAP_TYPE
 //
 //   MEDIAPIPE_REGISTER_TYPE(
-//       std::string, "string", StringSerializeFn, StringDeserializeFn);
+//       std::string, "std::string", StringSerializeFn, StringDeserializeFn);
 //
 #define MEDIAPIPE_REGISTER_TYPE(type, type_name, serialize_fn, deserialize_fn) \
   SET_MEDIAPIPE_TYPE_MAP_VALUE(                                                \
@@ -293,7 +293,7 @@ DEFINE_MEDIAPIPE_TYPE_MAP(PacketTypeStringToMediaPipeTypeData, std::string);
 //     typedef is used, the name should be prefixed with the namespace(s),
 //     seperated by double colons.
 //
-// Example 1: register type with non-std::string proxy.
+// Example 1: register type with non-string proxy.
 //   absl::Status ToProxyFn(
 //       const ClassType& obj, ProxyType* proxy)
 //   {
@@ -315,15 +315,15 @@ DEFINE_MEDIAPIPE_TYPE_MAP(PacketTypeStringToMediaPipeTypeData, std::string);
 //      ::mediapipe::DeserializeUsingGenericFn<ClassType WITH_MEDIAPIPE_PROXY
 //      ProxyType>, ToProxyFn, FromProxyFn);
 //
-// Example 2: register type with std::string proxy.
-//   absl::Status ToProxyFn(const ClassType& obj, std::string* encoding)
+// Example 2: register type with string proxy.
+//   absl::Status ToProxyFn(const ClassType& obj, string* encoding)
 //   {
 //     ...
 //     return absl::OkStatus();
 //   }
 //
 //   absl::Status FromProxyFn(
-//       const ProxyType& proxy, std::string* encoding) {
+//       const ProxyType& proxy, string* encoding) {
 //     ...
 //     return absl::OkStatus();
 //   }
@@ -367,7 +367,7 @@ inline const std::string* MediaPipeTypeStringFromTypeId(const size_t type_id) {
   return (value) ? &value->type_string : nullptr;
 }
 
-// Returns std::string identifier of type or NULL if not registered.
+// Returns string identifier of type or NULL if not registered.
 template <typename T>
 inline const std::string* MediaPipeTypeString() {
   return MediaPipeTypeStringFromTypeId(tool::GetTypeHash<T>());

mediapipe/framework/validated_graph_config.cc

@@ -49,7 +49,7 @@ namespace mediapipe {
 
 namespace {
 
-// Create a debug std::string name for a set of edge.  An edge can be either
+// Create a debug string name for a set of edge.  An edge can be either
 // a stream or a side packet.
 std::string DebugEdgeNames(
     const std::string& edge_type,
@@ -413,8 +413,11 @@ absl::Status ValidatedGraphConfig::Initialize(
 
   // Set Any types based on what they connect to.
   MP_RETURN_IF_ERROR(ResolveAnyTypes(&input_streams_, &output_streams_));
+  MP_RETURN_IF_ERROR(ResolveOneOfTypes(&input_streams_, &output_streams_));
   MP_RETURN_IF_ERROR(
       ResolveAnyTypes(&input_side_packets_, &output_side_packets_));
+  MP_RETURN_IF_ERROR(
+      ResolveOneOfTypes(&input_side_packets_, &output_side_packets_));
 
   // Validate consistency of side packets and streams.
   MP_RETURN_IF_ERROR(ValidateSidePacketTypes());
@@ -908,6 +911,29 @@ absl::Status ValidatedGraphConfig::ResolveAnyTypes(
   return absl::OkStatus();
 }
 
+absl::Status ValidatedGraphConfig::ResolveOneOfTypes(
+    std::vector<EdgeInfo>* input_edges, std::vector<EdgeInfo>* output_edges) {
+  for (EdgeInfo& input_edge : *input_edges) {
+    if (input_edge.upstream == -1) {
+      continue;
+    }
+    EdgeInfo& output_edge = (*output_edges)[input_edge.upstream];
+    PacketType* input_root = input_edge.packet_type->GetSameAs();
+    PacketType* output_root = output_edge.packet_type->GetSameAs();
+    if (!input_root->IsConsistentWith(*output_root)) continue;
+    // We narrow down OneOf types here if the other side is a single type.
+    // We do not currently intersect multiple OneOf types.
+    // Note that this is sensitive to the order edges are examined.
+    // TODO: we should be more thorough.
+    if (input_root->IsOneOf() && output_root->IsExactType()) {
+      input_root->SetSameAs(output_edge.packet_type);
+    } else if (output_root->IsOneOf() && input_root->IsExactType()) {
+      output_root->SetSameAs(input_edge.packet_type);
+    }
+  }
+  return absl::OkStatus();
+}
+
 absl::Status ValidatedGraphConfig::ValidateStreamTypes() {
   for (const EdgeInfo& stream : input_streams_) {
     RET_CHECK_NE(stream.upstream, -1);

mediapipe/framework/validated_graph_config.h

@@ -133,12 +133,11 @@ class NodeTypeInfo {
   // This function is only valid for a NodeTypeInfo of NodeType CALCULATOR.
   bool AddSource(int index) { return ancestor_sources_.insert(index).second; }
 
-  // Convert the NodeType enum into a std::string (generally for error
-  // messaging).
+  // Convert the NodeType enum into a string (generally for error messaging).
   static std::string NodeTypeToString(NodeType node_type);
 
-  // Returns the name of the specified InputStreamHandler, or empty std::string
-  // if none set.
+  // Returns the name of the specified InputStreamHandler, or empty string if
+  // none set.
   std::string GetInputStreamHandler() const {
     return contract_.GetInputStreamHandler();
   }
@@ -383,6 +382,9 @@ class ValidatedGraphConfig {
   // Infer the type of types set to "Any" by what they are connected to.
   absl::Status ResolveAnyTypes(std::vector<EdgeInfo>* input_edges,
                                std::vector<EdgeInfo>* output_edges);
+  // Narrow down OneOf types if they other end is a single type.
+  absl::Status ResolveOneOfTypes(std::vector<EdgeInfo>* input_edges,
+                                 std::vector<EdgeInfo>* output_edges);
 
   // Returns an error if the generator graph does not have consistent
   // type specifications for side packets.