Project import generated by Copybara.

GitOrigin-RevId: 27c70b5fe62ab71189d358ca122ee4b19c817a8f
2021-07-27 19:36:32 -04:00
parent 374f5e2e7e
commit 50c92c6623
158 changed files with 4704 additions and 621 deletions
@@ -28,6 +28,7 @@ package(default_visibility = ["//visibility:private"])

 exports_files([
    "simple_subgraph_template.cc",
+    "options_lib_template.cc",
 ])

 bzl_library(
@@ -40,6 +41,7 @@ bzl_library(
        ":build_defs_bzl",
        "//mediapipe/framework:encode_binary_proto",
        "//mediapipe/framework:transitive_protos_bzl",
+        "//mediapipe/framework/deps:descriptor_set_bzl",
        "//mediapipe/framework/deps:expand_template_bzl",
    ],
 )
@@ -71,6 +73,7 @@ cc_library(
 mediapipe_proto_library(
    name = "calculator_graph_template_proto",
    srcs = ["calculator_graph_template.proto"],
+    def_options_lib = False,
    def_py_proto = False,
    visibility = ["//visibility:public"],
    deps = [
@@ -80,6 +83,17 @@ mediapipe_proto_library(
    ],
 )

+mediapipe_proto_library(
+    name = "packet_generator_wrapper_calculator_proto",
+    srcs = ["packet_generator_wrapper_calculator.proto"],
+    def_py_proto = False,
+    visibility = ["//mediapipe/framework:mediapipe_internal"],
+    deps = [
+        "//mediapipe/framework:calculator_options_proto",
+        "//mediapipe/framework:packet_generator_proto",
+    ],
+)
+
 mediapipe_proto_library(
    name = "source_proto",
    srcs = ["source.proto"],
@@ -159,10 +173,10 @@ cc_library(
    hdrs = ["options_map.h"],
    visibility = ["//mediapipe/framework:mediapipe_internal"],
    deps = [
+        ":type_util",
        "//mediapipe/framework:calculator_cc_proto",
        "//mediapipe/framework/port:any_proto",
        "//mediapipe/framework/port:status",
-        "//mediapipe/framework/tool:type_util",
    ],
 )

@@ -190,13 +204,41 @@ cc_library(
    ],
 )

+cc_binary(
+    name = "message_type_util",
+    srcs = ["message_type_util.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//mediapipe/framework/port:advanced_proto",
+        "//mediapipe/framework/port:file_helpers",
+        "//mediapipe/framework/port:logging",
+        "@com_google_absl//absl/flags:flag",
+        "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "options_registry",
+    srcs = ["options_registry.cc"],
+    hdrs = ["options_registry.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//mediapipe/framework/deps:registration",
+        "//mediapipe/framework/port:advanced_proto",
+        "//mediapipe/framework/port:logging",
+    ],
+)
+
 mediapipe_cc_test(
    name = "options_util_test",
    size = "small",
    srcs = ["options_util_test.cc"],
+    # A non-empty "data" param is needed to build the "_test_wasm" target.
    data = [":node_chain_subgraph.proto"],
    requires_full_emulation = False,
    deps = [
+        ":options_registry",
        ":options_util",
        "//mediapipe/calculators/core:flow_limiter_calculator",
        "//mediapipe/calculators/core:flow_limiter_calculator_cc_proto",
@@ -207,11 +249,26 @@ mediapipe_cc_test(
        "//mediapipe/framework/port:gtest_main",
        "//mediapipe/framework/port:parse_text_proto",
        "//mediapipe/framework/port:status",
-        "//mediapipe/framework/testdata:night_light_calculator_cc_proto",
+        "//mediapipe/framework/testdata:night_light_calculator_options_lib",
+        "//mediapipe/framework/tool:node_chain_subgraph_options_lib",
        "//mediapipe/util:header_util",
    ],
 )

+cc_library(
+    name = "packet_generator_wrapper_calculator",
+    srcs = ["packet_generator_wrapper_calculator.cc"],
+    visibility = ["//mediapipe/framework:mediapipe_internal"],
+    deps = [
+        ":packet_generator_wrapper_calculator_cc_proto",
+        "//mediapipe/framework:calculator_base",
+        "//mediapipe/framework:calculator_registry",
+        "//mediapipe/framework:output_side_packet",
+        "//mediapipe/framework:packet_generator",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
    name = "proto_util_lite",
    srcs = ["proto_util_lite.cc"],
@@ -776,6 +833,7 @@ cc_test(
        "//mediapipe/framework/port:ret_check",
        "//mediapipe/framework/port:status",
        "//mediapipe/framework/stream_handler:immediate_input_stream_handler",
+        "@com_google_absl//absl/strings",
    ],
 )

@@ -19,6 +19,7 @@ load("//mediapipe/framework:encode_binary_proto.bzl", "encode_binary_proto", "ge
 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")

 def mediapipe_binary_graph(name, graph = None, output_name = None, deps = [], testonly = False, **kwargs):
    """Converts a graph from text format to binary format."""
@@ -152,3 +153,115 @@ def mediapipe_simple_subgraph(
        testonly = testonly,
        **kwargs
    )
+
+def mediapipe_reexport_library(
+        name,
+        actual,
+        **kwargs):
+    """Defines a cc_library that exports the headers of other libraries.
+
+    Normally cc_library does not export the headers of its dependencies,
+    and the clang "layering_check" requires clients to depend on them
+    directly.  Header files can be exported by listing them in either
+    cc_library's "hdrs" or "textual_hdrs" argument.  The "textual_hdrs"
+    argument can also accept library targets and has the effect of
+    exporting their header files and permitting client references to them.
+    The result is a new library target that combines and exports the public
+    interfaces of several existing library targets.
+
+    Args:
+      name: the name for the combined target.
+      actual: the targets to combine and export together.
+      **kwargs: Remaining keyword args, forwarded to cc_library.
+    """
+    native.cc_library(
+        name = name,
+        textual_hdrs = actual,
+        deps = actual,
+        **kwargs
+    )
+
+def mediapipe_options_library(
+        name,
+        proto_lib,
+        deps = [],
+        visibility = None,
+        testonly = None,
+        **kwargs):
+    """Registers options protobuf metadata for defining options packets.
+
+    Args:
+      name: name of the options_lib target to define.
+      proto_lib: the proto_library target to register.
+      deps: any additional protobuf dependencies.
+      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.
+    """
+
+    transitive_descriptor_set(
+        name = proto_lib + "_transitive",
+        deps = [proto_lib],
+        testonly = testonly,
+    )
+    direct_descriptor_set(
+        name = proto_lib + "_direct",
+        deps = [proto_lib],
+        testonly = testonly,
+    )
+    data_as_c_string(
+        name = name + "_inc",
+        srcs = [proto_lib + "_transitive-transitive-descriptor-set.proto.bin"],
+        outs = [proto_lib + "_descriptors.inc"],
+    )
+    native.genrule(
+        name = name + "_type_name",
+        srcs = [proto_lib + "_direct-direct-descriptor-set.proto.bin"],
+        outs = [name + "_type_name.h"],
+        cmd = ("$(location " + "//mediapipe/framework/tool:message_type_util" + ") " +
+               ("--input_path=$(location %s) " % (proto_lib + "_direct-direct-descriptor-set.proto.bin")) +
+               ("--root_type_macro_output_path=$(location %s) " % (name + "_type_name.h"))),
+        tools = ["//mediapipe/framework/tool:message_type_util"],
+        visibility = visibility,
+        testonly = testonly,
+    )
+    expand_template(
+        name = name + "_cc",
+        template = clean_dep("//mediapipe/framework/tool:options_lib_template.cc"),
+        out = name + ".cc",
+        substitutions = {
+            "{{MESSAGE_NAME_HEADER}}": native.package_name() + "/" + name + "_type_name.h",
+            "{{MESSAGE_PROTO_HEADER}}": native.package_name() + "/" + proto_lib.replace("_proto", ".pb.h"),
+            "{{DESCRIPTOR_INC_FILE_PATH}}": native.package_name() + "/" + proto_lib + "_descriptors.inc",
+        },
+        testonly = testonly,
+    )
+    native.cc_library(
+        name = proto_lib.replace("_proto", "_options_registry"),
+        srcs = [
+            name + ".cc",
+            proto_lib + "_descriptors.inc",
+            name + "_type_name.h",
+        ],
+        deps = [
+            clean_dep("//mediapipe/framework:calculator_framework"),
+            clean_dep("//mediapipe/framework/port:advanced_proto"),
+            clean_dep("//mediapipe/framework/tool:options_registry"),
+            proto_lib.replace("_proto", "_cc_proto"),
+        ] + deps,
+        alwayslink = 1,
+        visibility = visibility,
+        testonly = testonly,
+        features = ["-no_undefined"],
+        **kwargs
+    )
+    mediapipe_reexport_library(
+        name = name,
+        actual = [
+            proto_lib.replace("_proto", "_cc_proto"),
+            proto_lib.replace("_proto", "_options_registry"),
+        ],
+        visibility = visibility,
+        testonly = testonly,
+        **kwargs
+    )
@@ -0,0 +1,170 @@
+
+#include <iostream>
+#include <string>
+
+#include "absl/flags/flag.h"
+#include "absl/flags/parse.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_replace.h"
+#include "mediapipe/framework/port/advanced_proto_inc.h"
+#include "mediapipe/framework/port/file_helpers.h"
+#include "mediapipe/framework/port/logging.h"
+
+ABSL_FLAG(std::string, input_path, "",
+          "Full path of the FileDescriptorSet to summarize. ");
+ABSL_FLAG(std::string, root_type_name_output_path, "",
+          "Where to write the output root message type name. ");
+ABSL_FLAG(std::string, root_type_macro_output_path, "",
+          "Where to write the output root message type macro. ");
+
+namespace mediapipe {
+
+using proto_ns::DescriptorProto;
+using proto_ns::FileDescriptorProto;
+using proto_ns::FileDescriptorSet;
+
+// Utility to extract summary data about protobuf descriptors.
+//
+// This utility is currently used by the build rule mediapipe_options_library()
+// to recover the package-name and type-name associated with each
+// mediapipe_proto_library() target.
+class DescriptorReader {
+ public:
+  // Returns a FileDescriptor that is not referenced by other FileDescriptors
+  // in a FileDescriptorSet.
+  static FileDescriptorProto FindTopFile(const FileDescriptorSet& files) {
+    std::set<std::string> file_names;
+    for (const FileDescriptorProto& file : files.file()) {
+      file_names.insert(file.name());
+    }
+    for (const FileDescriptorProto& file : files.file()) {
+      for (const std::string& dep : file.dependency()) {
+        file_names.erase(dep);
+      }
+    }
+    for (const FileDescriptorProto& file : files.file()) {
+      if (file_names.count(file.name()) > 0) {
+        return file;
+      }
+    }
+    return FileDescriptorProto();
+  }
+
+  static std::string CleanTypeName(const std::string& type_name) {
+    return (type_name.rfind('.', 0) == 0) ? type_name.substr(1) : type_name;
+  }
+
+  static std::string CleanTypeName(const std::string& package,
+                                   const std::string& name) {
+    return absl::StrCat(package, ".", name);
+  }
+
+  // Returns the length of the common prefix between two strings.
+  static int MatchingPrefixLength(const std::string& s, const std::string& t) {
+    int i = 0;
+    while (i < std::min(s.size(), t.size()) && s[i] == t[i]) {
+      ++i;
+    }
+    return i;
+  }
+
+  // Returns the type-name that best matches the descriptor file-name.
+  static std::string BestTypeName(const std::set<std::string>& type_names,
+                                  const FileDescriptorProto& file) {
+    std::string proto_name = std::string(file::Basename(file.name()));
+    proto_name = proto_name.substr(
+        0, proto_name.size() - file::Extension(proto_name).size() - 1);
+    proto_name.erase(std::remove(proto_name.begin(), proto_name.end(), '_'),
+                     proto_name.end());
+    std::string result = "";
+    int best_match = -1;
+    for (const std::string& type_name : type_names) {
+      std::string name = absl::AsciiStrToLower(type_name);
+      if (name.rfind('.') != std::string::npos) {
+        name = name.substr(name.rfind('.') + 1);
+      }
+      int m = MatchingPrefixLength(proto_name, name);
+      if (m > best_match) {
+        best_match = m;
+        result = type_name;
+      }
+    }
+    return result;
+  }
+
+  // Returns a DescriptorProto that is not referenced by other DescriptorProtos
+  // in a FileDescriptorProto.
+  static DescriptorProto FindTopDescriptor(const FileDescriptorProto& file) {
+    std::set<std::string> type_names;
+    std::set<std::string> refs;
+    for (const DescriptorProto& descriptor : file.message_type()) {
+      type_names.insert(CleanTypeName(file.package(), descriptor.name()));
+    }
+    std::string best_name = BestTypeName(type_names, file);
+    for (const DescriptorProto& descriptor : file.message_type()) {
+      if (best_name == CleanTypeName(file.package(), descriptor.name())) {
+        return descriptor;
+      }
+    }
+    return DescriptorProto();
+  }
+
+  static std::string FindTopTypeName(const FileDescriptorSet& files) {
+    FileDescriptorProto file = FindTopFile(files);
+    DescriptorProto descriptor = FindTopDescriptor(file);
+    return CleanTypeName(file.package(), descriptor.name());
+  }
+
+  static FileDescriptorSet ReadFileDescriptorSet(const std::string& path) {
+    std::string contents;
+    CHECK_OK(file::GetContents(path, &contents));
+    proto_ns::FileDescriptorSet result;
+    result.ParseFromString(contents);
+    return result;
+  }
+
+  static void WriteFile(const std::string& path, const std::string& contents) {
+    CHECK_OK(file::SetContents(path, contents));
+  }
+
+  static void WriteMessageTypeName(const std::string& path,
+                                   const FileDescriptorSet& files) {
+    FileDescriptorProto file = FindTopFile(files);
+    DescriptorProto descriptor = FindTopDescriptor(file);
+    std::string type_name = mediapipe::DescriptorReader::FindTopTypeName(files);
+    mediapipe::DescriptorReader::WriteFile(
+        absl::GetFlag(FLAGS_root_type_name_output_path), type_name);
+  }
+
+  static void WriteMessageTypeMacro(const std::string& path,
+                                    const FileDescriptorSet& files) {
+    FileDescriptorProto file = FindTopFile(files);
+    DescriptorProto descriptor = FindTopDescriptor(file);
+    std::string type_package =
+        absl::StrReplaceAll(file.package(), {{".", "::"}});
+    std::string type_name = descriptor.name();
+    std::string contents =
+        absl::StrCat("#define MP_OPTION_TYPE_NS ", type_package, "\n") +
+        absl::StrCat("#define MP_OPTION_TYPE_NAME ", type_name, "\n");
+    WriteFile(path, contents);
+  }
+};
+
+}  // namespace mediapipe
+
+int main(int argc, char** argv) {
+  google::InitGoogleLogging(argv[0]);
+  absl::ParseCommandLine(argc, argv);
+  auto files = mediapipe::DescriptorReader::ReadFileDescriptorSet(
+      absl::GetFlag(FLAGS_input_path));
+  if (!absl::GetFlag(FLAGS_root_type_name_output_path).empty()) {
+    mediapipe::DescriptorReader::WriteMessageTypeName(
+        absl::GetFlag(FLAGS_root_type_name_output_path), files);
+  }
+  if (!absl::GetFlag(FLAGS_root_type_macro_output_path).empty()) {
+    mediapipe::DescriptorReader::WriteMessageTypeMacro(
+        absl::GetFlag(FLAGS_root_type_macro_output_path), files);
+  }
+  return EXIT_SUCCESS;
+}
@@ -0,0 +1,44 @@
+// Copyright 2019 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// This template is used by the mediapipe_simple_subgraph macro in
+// //mediapipe/framework/tool/mediapipe_graph.bzl
+
+#include "mediapipe/framework/port/advanced_proto_inc.h"
+#include "mediapipe/framework/tool/options_registry.h"
+#include "{{MESSAGE_NAME_HEADER}}"
+#include "{{MESSAGE_PROTO_HEADER}}"
+
+namespace {
+constexpr char kDescriptorContents[] =
+#include "{{DESCRIPTOR_INC_FILE_PATH}}"
+    ;  // NOLINT(whitespace/semicolon)
+
+mediapipe::proto_ns::FileDescriptorSet ParseFileDescriptorSet(
+    const std::string& pb) {
+  mediapipe::proto_ns::FileDescriptorSet files;
+  files.ParseFromString(pb);
+  return files;
+}
+
+}  // namespace
+
+namespace mediapipe {
+// The protobuf descriptor for an options message type.
+template <>
+const RegistrationToken tool::OptionsRegistry::registration_token<
+    MP_OPTION_TYPE_NS::MP_OPTION_TYPE_NAME> =
+    tool::OptionsRegistry::Register(ParseFileDescriptorSet(
+        std::string(kDescriptorContents, sizeof(kDescriptorContents) - 1)));
+}  // namespace mediapipe
@@ -0,0 +1,47 @@
+#include "mediapipe/framework/tool/options_registry.h"
+
+namespace mediapipe {
+namespace tool {
+
+proto_ns::DescriptorPool* OptionsRegistry::options_descriptor_pool() {
+  static proto_ns::DescriptorPool* result = new proto_ns::DescriptorPool();
+  return result;
+}
+
+RegistrationToken OptionsRegistry::Register(
+    const proto_ns::FileDescriptorSet& files) {
+  for (auto& file : files.file()) {
+    options_descriptor_pool()->BuildFile(file);
+  }
+  return RegistrationToken([]() {});
+}
+
+const proto_ns::Descriptor* OptionsRegistry::GetProtobufDescriptor(
+    const std::string& type_name) {
+  const proto_ns::Descriptor* result =
+      proto_ns::DescriptorPool::generated_pool()->FindMessageTypeByName(
+          type_name);
+  if (!result) {
+    result = options_descriptor_pool()->FindMessageTypeByName(type_name);
+  }
+  return result;
+}
+
+void OptionsRegistry::FindAllExtensions(
+    const proto_ns::Descriptor& extendee,
+    std::vector<const proto_ns::FieldDescriptor*>* result) {
+  using proto_ns::DescriptorPool;
+  std::vector<const proto_ns::FieldDescriptor*> extensions;
+  DescriptorPool::generated_pool()->FindAllExtensions(&extendee, &extensions);
+  options_descriptor_pool()->FindAllExtensions(&extendee, &extensions);
+  absl::flat_hash_set<int> numbers;
+  for (const proto_ns::FieldDescriptor* extension : extensions) {
+    bool inserted = numbers.insert(extension->number()).second;
+    if (inserted) {
+      result->push_back(extension);
+    }
+  }
+}
+
+}  // namespace tool
+}  // namespace mediapipe
@@ -0,0 +1,40 @@
+#ifndef MEDIAPIPE_FRAMEWORK_TOOL_OPTIONS_REGISTRY_H_
+#define MEDIAPIPE_FRAMEWORK_TOOL_OPTIONS_REGISTRY_H_
+
+#include "mediapipe/framework/deps/registration.h"
+#include "mediapipe/framework/port/advanced_proto_inc.h"
+
+namespace mediapipe {
+namespace tool {
+
+// A static registry that stores descriptors for protobufs used in MediaPipe
+// calculator options. Lite-proto builds do not normally include descriptors.
+// These registered descriptors allow individual protobuf fields to be
+// referenced and specified separately within CalculatorGraphConfigs.
+class OptionsRegistry {
+ public:
+  // Registers the protobuf descriptors for a MessageLite.
+  static RegistrationToken Register(const proto_ns::FileDescriptorSet& files);
+
+  // Finds the descriptor for a protobuf.
+  static const proto_ns::Descriptor* GetProtobufDescriptor(
+      const std::string& type_name);
+
+  // Returns all known proto2 extensions to a type.
+  static void FindAllExtensions(
+      const proto_ns::Descriptor& extendee,
+      std::vector<const proto_ns::FieldDescriptor*>* result);
+
+ private:
+  // Stores the descriptors for each options protobuf type.
+  static proto_ns::DescriptorPool* options_descriptor_pool();
+
+  // Registers the descriptors for each options protobuf type.
+  template <class MessageT>
+  static const RegistrationToken registration_token;
+};
+
+}  // namespace tool
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_TOOL_OPTIONS_REGISTRY_H_
@@ -19,6 +19,8 @@
 #include "mediapipe/framework/port/gtest.h"
 #include "mediapipe/framework/port/parse_text_proto.h"
 #include "mediapipe/framework/port/status_matchers.h"
+#include "mediapipe/framework/testdata/night_light_calculator.pb.h"
+#include "mediapipe/framework/tool/options_registry.h"

 namespace mediapipe {
 namespace {
@@ -42,5 +44,13 @@ TEST_F(OptionsUtilTest, GetProtobufDescriptor) {
 #endif
 }

+// Retrieves the description of a protobuf from the OptionsRegistry.
+TEST_F(OptionsUtilTest, GetProtobufDescriptorRegistered) {
+  const proto_ns::Descriptor* descriptor =
+      tool::OptionsRegistry::GetProtobufDescriptor(
+          "mediapipe.NightLightCalculatorOptions");
+  EXPECT_NE(nullptr, descriptor);
+}
+
 }  // namespace
 }  // namespace mediapipe
@@ -0,0 +1,52 @@
+#include "mediapipe/framework/calculator_base.h"
+#include "mediapipe/framework/calculator_registry.h"
+#include "mediapipe/framework/output_side_packet.h"
+#include "mediapipe/framework/packet_generator.h"
+#include "mediapipe/framework/tool/packet_generator_wrapper_calculator.pb.h"
+
+namespace mediapipe {
+
+class PacketGeneratorWrapperCalculator : public CalculatorBase {
+ public:
+  static absl::Status GetContract(CalculatorContract* cc) {
+    const auto& options =
+        cc->Options<::mediapipe::PacketGeneratorWrapperCalculatorOptions>();
+    ASSIGN_OR_RETURN(auto static_access,
+                     mediapipe::internal::StaticAccessToGeneratorRegistry::
+                         CreateByNameInNamespace(options.package(),
+                                                 options.packet_generator()));
+    MP_RETURN_IF_ERROR(static_access->FillExpectations(
+                           options.options(), &cc->InputSidePackets(),
+                           &cc->OutputSidePackets()))
+            .SetPrepend()
+        << options.packet_generator() << "::FillExpectations() failed: ";
+    return absl::OkStatus();
+  }
+
+  absl::Status Open(CalculatorContext* cc) override {
+    const auto& options =
+        cc->Options<::mediapipe::PacketGeneratorWrapperCalculatorOptions>();
+    ASSIGN_OR_RETURN(auto static_access,
+                     mediapipe::internal::StaticAccessToGeneratorRegistry::
+                         CreateByNameInNamespace(options.package(),
+                                                 options.packet_generator()));
+    mediapipe::PacketSet output_packets(cc->OutputSidePackets().TagMap());
+    MP_RETURN_IF_ERROR(static_access->Generate(options.options(),
+                                               cc->InputSidePackets(),
+                                               &output_packets))
+            .SetPrepend()
+        << options.packet_generator() << "::Generate() failed: ";
+    for (auto id = output_packets.BeginId(); id < output_packets.EndId();
+         ++id) {
+      cc->OutputSidePackets().Get(id).Set(output_packets.Get(id));
+    }
+    return absl::OkStatus();
+  }
+
+  absl::Status Process(CalculatorContext* cc) override {
+    return absl::OkStatus();
+  }
+};
+REGISTER_CALCULATOR(PacketGeneratorWrapperCalculator);
+
+}  // namespace mediapipe
@@ -0,0 +1,19 @@
+syntax = "proto2";
+
+package mediapipe;
+
+import "mediapipe/framework/calculator_options.proto";
+import "mediapipe/framework/packet_generator.proto";
+
+message PacketGeneratorWrapperCalculatorOptions {
+  extend CalculatorOptions {
+    optional PacketGeneratorWrapperCalculatorOptions ext = 381945445;
+  }
+
+  // Same as the corresponding fields in PacketGeneratorConfig.
+  optional string packet_generator = 1;
+  optional PacketGeneratorOptions options = 2;
+  // Same as CalculatorGraphConfig.package. Copied here since the graph config
+  // is not available to the calculator.
+  optional string package = 3;
+}
@@ -186,7 +186,6 @@ absl::Status FindCorrespondingStreams(
 absl::Status ValidateSubgraphFields(
    const CalculatorGraphConfig::Node& subgraph_node) {
  if (subgraph_node.source_layer() || subgraph_node.buffer_size_hint() ||
-      subgraph_node.has_input_stream_handler() ||
      subgraph_node.has_output_stream_handler() ||
      subgraph_node.input_stream_info_size() != 0 ||
      !subgraph_node.executor().empty()) {
@@ -59,6 +59,11 @@ using mediapipe::SwitchContainerOptions;
 // or contained_node 1, given "ENABLE:false" or "ENABLE:true" respectively.
 // Input-side-packet "ENABLE" and input-stream "SELECT" can also be used
 // similarly to specify the active channel.
+//
+// Note that this container defaults to use ImmediateInputStreamHandler,
+// which can be used to accept infrequent "enable" packets asynchronously.
+// However, it can be overridden to work with DefaultInputStreamHandler,
+// which can be used to accept frequent "enable" packets synchronously.
 class SwitchContainer : public Subgraph {
 public:
  SwitchContainer() = default;
@@ -79,11 +84,16 @@ std::string ChannelName(const std::string& name, int channel) {
 // Returns a SwitchDemuxCalculator node.
 CalculatorGraphConfig::Node* BuildDemuxNode(
    const std::map<TagIndex, std::string>& input_tags,
+    const CalculatorGraphConfig::Node& container_node,
    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;
 }

@@ -233,7 +243,7 @@ absl::StatusOr<CalculatorGraphConfig> SwitchContainer::GetConfig(
  ParseTags(container_streams.output_side_packet(), &side_output_tags);

  // Add a graph node for the demux, mux.
-  auto demux = BuildDemuxNode(input_tags, &config);
+  auto demux = BuildDemuxNode(input_tags, container_node, &config);
  CopyOptions(container_node, demux);
  ClearContainerOptions(demux);
  demux->add_input_stream("SELECT:gate_select");
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

+#include "absl/strings/str_replace.h"
 #include "mediapipe/framework/calculator.pb.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/deps/message_matchers.h"
@@ -66,8 +67,9 @@ 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() {
-  return mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+CalculatorGraphConfig SubnodeContainerExample(
+    const std::string& input_stream_handler = "") {
+  std::string config = R"pb(
    input_stream: "foo"
    input_stream: "enable"
    input_side_packet: "timezone"
@@ -81,7 +83,7 @@ CalculatorGraphConfig SubnodeContainerExample() {
          contained_node: { calculator: "TripleIntCalculator" }
          contained_node: { calculator: "PassThroughCalculator" }
        }
-      }
+      } $input_stream_handler
    }
    node {
      calculator: "PassThroughCalculator"
@@ -90,7 +92,11 @@ CalculatorGraphConfig SubnodeContainerExample() {
      output_stream: "output_foo"
      output_stream: "output_bar"
    }
-  )pb");
+  )pb";
+
+  return mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
+      absl::StrReplaceAll(config,
+                          {{"$input_stream_handler", input_stream_handler}}));
 }

 // A testing example of a SwitchContainer containing two subnodes.
@@ -124,7 +130,8 @@ CalculatorGraphConfig SideSubnodeContainerExample() {
 }

 // Runs the test container graph with a few input packets.
-void RunTestContainer(CalculatorGraphConfig supergraph) {
+void RunTestContainer(CalculatorGraphConfig supergraph,
+                      bool send_bounds = false) {
  CalculatorGraph graph;
  std::vector<Packet> out_foo, out_bar;
  tool::AddVectorSink("output_foo", &supergraph, &out_foo);
@@ -132,17 +139,23 @@ void RunTestContainer(CalculatorGraphConfig supergraph) {
  MP_ASSERT_OK(graph.Initialize(supergraph, {}));
  MP_ASSERT_OK(graph.StartRun({{"timezone", MakePacket<int>(3)}}));

-  // Send enable == true signal at 5000 us.
-  const int64 enable_ts = 5000;
-  MP_EXPECT_OK(graph.AddPacketToInputStream(
-      "enable", MakePacket<bool>(true).At(Timestamp(enable_ts))));
-  MP_ASSERT_OK(graph.WaitUntilIdle());
+  if (!send_bounds) {
+    // Send enable == true signal at 5000 us.
+    const int64 enable_ts = 5000;
+    MP_EXPECT_OK(graph.AddPacketToInputStream(
+        "enable", MakePacket<bool>(true).At(Timestamp(enable_ts))));
+    MP_ASSERT_OK(graph.WaitUntilIdle());
+  }

  const int packet_count = 10;
  // Send int value packets at {10K, 20K, 30K, ..., 100K}.
  for (uint64 t = 1; t <= packet_count; ++t) {
    MP_EXPECT_OK(graph.AddPacketToInputStream(
        "foo", MakePacket<int>(t).At(Timestamp(t * 10000))));
+    if (send_bounds) {
+      MP_EXPECT_OK(graph.AddPacketToInputStream(
+          "enable", MakePacket<bool>(true).At(Timestamp(t * 10000))));
+    }
    MP_ASSERT_OK(graph.WaitUntilIdle());
    // The inputs are sent to the input stream "foo", they should pass through.
    EXPECT_EQ(out_foo.size(), t);
@@ -153,15 +166,21 @@ void RunTestContainer(CalculatorGraphConfig supergraph) {
    EXPECT_EQ(out_bar.back().Get<int>(), t);
  }

-  // Send enable == false signal at 105K us.
-  MP_EXPECT_OK(graph.AddPacketToInputStream(
-      "enable", MakePacket<bool>(false).At(Timestamp(105000))));
-  MP_ASSERT_OK(graph.WaitUntilIdle());
+  if (!send_bounds) {
+    // Send enable == false signal at 105K us.
+    MP_EXPECT_OK(graph.AddPacketToInputStream(
+        "enable", MakePacket<bool>(false).At(Timestamp(105000))));
+    MP_ASSERT_OK(graph.WaitUntilIdle());
+  }

  // Send int value packets at {110K, 120K, ..., 200K}.
  for (uint64 t = 11; t <= packet_count * 2; ++t) {
    MP_EXPECT_OK(graph.AddPacketToInputStream(
        "foo", MakePacket<int>(t).At(Timestamp(t * 10000))));
+    if (send_bounds) {
+      MP_EXPECT_OK(graph.AddPacketToInputStream(
+          "enable", MakePacket<bool>(false).At(Timestamp(t * 10000))));
+    }
    MP_ASSERT_OK(graph.WaitUntilIdle());
    // The inputs are sent to the input stream "foo", they should pass through.
    EXPECT_EQ(out_foo.size(), t);
@@ -351,6 +370,86 @@ 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");
+  MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
+  RunTestContainer(supergraph, true);
+}
+
 // Shows the SwitchContainer container applied to a pair of simple subnodes.
 TEST(SwitchContainerTest, ApplyToSideSubnodes) {
  EXPECT_TRUE(SubgraphRegistry::IsRegistered("SwitchContainer"));
@@ -53,7 +53,7 @@ absl::Status RunGeneratorFillExpectations(
        << " is not a registered packet generator.");

  CalculatorContract contract;
-  MP_RETURN_IF_ERROR(contract.Initialize(config));
+  MP_RETURN_IF_ERROR(contract.Initialize(config, ""));

  {
    LegacyCalculatorSupport::Scoped<CalculatorContract> s(&contract);