Project import generated by Copybara.

PiperOrigin-RevId: 253489161
2019-06-16 16:06:57 -07:00
commit d68f5e4169
844 changed files with 134997 additions and 0 deletions
@@ -0,0 +1,118 @@
+// 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.
+//
+// Facility for registering subgraphs that can be included in other graphs.
+
+#ifndef MEDIAPIPE_FRAMEWORK_SUBGRAPH_H_
+#define MEDIAPIPE_FRAMEWORK_SUBGRAPH_H_
+
+#include "absl/base/macros.h"
+#include "absl/memory/memory.h"
+#include "mediapipe/framework/calculator.pb.h"
+#include "mediapipe/framework/deps/registration.h"
+#include "mediapipe/framework/port/status.h"
+#include "mediapipe/framework/port/statusor.h"
+#include "mediapipe/framework/tool/calculator_graph_template.pb.h"
+
+namespace mediapipe {
+
+// Instances of this class are responsible for providing a subgraph config.
+// They are only used during graph construction. They do not stay alive once
+// the graph is running.
+class Subgraph {
+ public:
+  using SubgraphOptions = CalculatorOptions;
+  Subgraph();
+  virtual ~Subgraph();
+  // Returns the config to use for one instantiation of the subgraph. The
+  // nodes and generators in this config will replace the subgraph node in
+  // the parent graph.
+  // Subclasses may use the options argument to parameterize the config.
+  // TODO: make this static?
+  virtual ::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
+      const SubgraphOptions& options) = 0;
+};
+
+using SubgraphRegistry = GlobalFactoryRegistry<std::unique_ptr<Subgraph>>;
+
+#define REGISTER_MEDIAPIPE_GRAPH(name)                               \
+  REGISTER_FACTORY_FUNCTION_QUALIFIED(::mediapipe::SubgraphRegistry, \
+                                      subgraph_registration, name,   \
+                                      absl::make_unique<name>)
+
+// A graph factory holding a literal CalculatorGraphConfig.
+class ProtoSubgraph : public Subgraph {
+ public:
+  ProtoSubgraph(const CalculatorGraphConfig& config);
+  virtual ~ProtoSubgraph();
+  virtual ::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
+      const Subgraph::SubgraphOptions& options);
+
+ private:
+  CalculatorGraphConfig config_;
+};
+
+// A graph factory holding a literal CalculatorGraphTemplate.
+class TemplateSubgraph : public Subgraph {
+ public:
+  TemplateSubgraph(const CalculatorGraphTemplate& templ);
+  virtual ~TemplateSubgraph();
+  virtual ::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
+      const Subgraph::SubgraphOptions& options);
+
+ private:
+  CalculatorGraphTemplate templ_;
+};
+
+// A local registry of CalculatorGraphConfig definitions.
+class GraphRegistry {
+ public:
+  // Creates a GraphRegistry derived from the GlobalFactoryRegistry,
+  // which stores statically linked Subgraphs.
+  GraphRegistry();
+
+  // Creates a GraphRegistry derived from the specified FunctionRegistry,
+  // which is used in place of the GlobalFactoryRegistry.
+  // Ownership of the specified FunctionRegistry is not transferred.
+  GraphRegistry(FunctionRegistry<std::unique_ptr<Subgraph>>* factories);
+
+  // Registers a graph config by name.
+  void Register(const std::string& type_name,
+                const CalculatorGraphConfig& config);
+
+  // Registers a template graph config by name.
+  void Register(const std::string& type_name,
+                const CalculatorGraphTemplate& templ);
+
+  // Returns true if the specified graph config is registered.
+  bool IsRegistered(const std::string& ns, const std::string& type_name) const;
+
+  // Returns the specified graph config.
+  ::mediapipe::StatusOr<CalculatorGraphConfig> CreateByName(
+      const std::string& ns, const std::string& type_name,
+      const Subgraph::SubgraphOptions* options = nullptr) const;
+
+  static GraphRegistry global_graph_registry;
+
+ private:
+  // The FunctionRegistry for dynamically loaded Subgraphs.
+  mutable FunctionRegistry<std::unique_ptr<Subgraph>> local_factories_;
+  // The FunctionRegistry for statically linked Subgraphs.
+  // The global_factories_ registry is overridden by local_factories_.
+  mutable FunctionRegistry<std::unique_ptr<Subgraph>>* global_factories_;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_SUBGRAPH_H_