diff --git a/mediapipe/tasks/cc/core/model_task_graph.cc b/mediapipe/tasks/cc/core/model_task_graph.cc index 47334b67..66434483 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.cc +++ b/mediapipe/tasks/cc/core/model_task_graph.cc @@ -156,21 +156,24 @@ absl::StatusOr ModelTaskGraph::GetConfig( } absl::StatusOr ModelTaskGraph::CreateModelResources( - SubgraphContext* sc, std::unique_ptr external_file) { + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix) { auto model_resources_cache_service = sc->Service(kModelResourcesCacheService); if (!model_resources_cache_service.IsAvailable()) { - ASSIGN_OR_RETURN(local_model_resources_, + ASSIGN_OR_RETURN(auto local_model_resource, ModelResources::Create("", std::move(external_file))); LOG(WARNING) << "A local ModelResources object is created. Please consider using " "ModelResourcesCacheService to cache the created ModelResources " "object in the CalculatorGraph."; - return local_model_resources_.get(); + local_model_resources_.push_back(std::move(local_model_resource)); + return local_model_resources_.back().get(); } ASSIGN_OR_RETURN( auto op_resolver_packet, model_resources_cache_service.GetObject().GetGraphOpResolverPacket()); - const std::string tag = CreateModelResourcesTag(sc->OriginalNode()); + const std::string tag = + absl::StrCat(CreateModelResourcesTag(sc->OriginalNode()), tag_suffix); ASSIGN_OR_RETURN(auto model_resources, ModelResources::Create(tag, std::move(external_file), op_resolver_packet)); @@ -182,7 +185,8 @@ absl::StatusOr ModelTaskGraph::CreateModelResources( absl::StatusOr ModelTaskGraph::CreateModelAssetBundleResources( - SubgraphContext* sc, std::unique_ptr external_file) { + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix) { auto model_resources_cache_service = sc->Service(kModelResourcesCacheService); bool has_file_pointer_meta = external_file->has_file_pointer_meta(); // if external file is set by file pointer, no need to add the model asset @@ -190,7 +194,7 @@ ModelTaskGraph::CreateModelAssetBundleResources( // not owned by this model asset bundle resources. if (!model_resources_cache_service.IsAvailable() || has_file_pointer_meta) { ASSIGN_OR_RETURN( - local_model_asset_bundle_resources_, + auto local_model_asset_bundle_resource, ModelAssetBundleResources::Create("", std::move(external_file))); if (!has_file_pointer_meta) { LOG(WARNING) @@ -198,10 +202,12 @@ ModelTaskGraph::CreateModelAssetBundleResources( "ModelResourcesCacheService to cache the created ModelResources " "object in the CalculatorGraph."; } - return local_model_asset_bundle_resources_.get(); + local_model_asset_bundle_resources_.push_back( + std::move(local_model_asset_bundle_resource)); + return local_model_asset_bundle_resources_.back().get(); } - const std::string tag = - CreateModelAssetBundleResourcesTag(sc->OriginalNode()); + const std::string tag = absl::StrCat( + CreateModelAssetBundleResourcesTag(sc->OriginalNode()), tag_suffix); ASSIGN_OR_RETURN( auto model_bundle_resources, ModelAssetBundleResources::Create(tag, std::move(external_file))); diff --git a/mediapipe/tasks/cc/core/model_task_graph.h b/mediapipe/tasks/cc/core/model_task_graph.h index 5ee70e8f..50dcc903 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.h +++ b/mediapipe/tasks/cc/core/model_task_graph.h @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include "absl/status/statusor.h" #include "absl/strings/str_format.h" @@ -75,9 +76,14 @@ class ModelTaskGraph : public Subgraph { // construction stage. Note that the external file contents will be moved // into the model resources object on creation. The returned model resources // pointer will provide graph authors with the access to the metadata - // extractor and the tflite model. + // extractor and the tflite model. When the model resources graph service is + // available, a tag is generated internally asscoiated with the created model + // resource. If more than one model resources are created in a graph, the + // model resources graph service add the tag_suffix to support multiple + // resources. absl::StatusOr CreateModelResources( - SubgraphContext* sc, std::unique_ptr external_file); + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix = ""); // If the model resources graph service is available, creates a model asset // bundle resources object from the subgraph context, and caches the created @@ -103,10 +109,15 @@ class ModelTaskGraph : public Subgraph { // that can only be used in the graph construction stage. Note that the // external file contents will be moved into the model asset bundle resources // object on creation. The returned model asset bundle resources pointer will - // provide graph authors with the access to extracted model files. + // provide graph authors with the access to extracted model files. When the + // model resources graph service is available, a tag is generated internally + // asscoiated with the created model asset bundle resource. If more than one + // model asset bundle resources are created in a graph, the model resources + // graph service add the tag_suffix to support multiple resources. absl::StatusOr CreateModelAssetBundleResources( - SubgraphContext* sc, std::unique_ptr external_file); + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix = ""); // Inserts a mediapipe task inference subgraph into the provided // GraphBuilder. The returned node provides the following interfaces to the @@ -124,9 +135,9 @@ class ModelTaskGraph : public Subgraph { api2::builder::Graph& graph) const; private: - std::unique_ptr local_model_resources_; + std::vector> local_model_resources_; - std::unique_ptr + std::vector> local_model_asset_bundle_resources_; };