diff --git a/mediapipe/tasks/cc/core/model_task_graph.cc b/mediapipe/tasks/cc/core/model_task_graph.cc index 46cc088f2..225fad418 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.cc +++ b/mediapipe/tasks/cc/core/model_task_graph.cc @@ -186,6 +186,21 @@ absl::StatusOr ModelTaskGraph::CreateModelResources( return model_resources_cache_service.GetObject().GetModelResources(tag); } +absl::StatusOr ModelTaskGraph::GetOrCreateModelResources( + SubgraphContext* sc, std::unique_ptr external_file, + std::string tag_suffix) { + auto model_resources_cache_service = sc->Service(kModelResourcesCacheService); + if (model_resources_cache_service.IsAvailable()) { + std::string tag = + absl::StrCat(CreateModelResourcesTag(sc->OriginalNode()), tag_suffix); + if (model_resources_cache_service.GetObject().Exists(tag)) { + return model_resources_cache_service.GetObject().GetModelResources(tag); + } + } + return ModelTaskGraph::CreateModelResources(sc, std::move(external_file), + tag_suffix); +} + absl::StatusOr ModelTaskGraph::CreateModelAssetBundleResources( SubgraphContext* sc, std::unique_ptr external_file, diff --git a/mediapipe/tasks/cc/core/model_task_graph.h b/mediapipe/tasks/cc/core/model_task_graph.h index 10634d3d0..38367da8f 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.h +++ b/mediapipe/tasks/cc/core/model_task_graph.h @@ -87,6 +87,20 @@ class ModelTaskGraph : public Subgraph { SubgraphContext* sc, std::unique_ptr external_file, std::string tag_suffix = ""); + template + absl::StatusOr GetOrCreateModelResources( + SubgraphContext* sc, std::string tag_suffix = "") { + auto external_file = std::make_unique(); + external_file->Swap(sc->MutableOptions() + ->mutable_base_options() + ->mutable_model_asset()); + return GetOrCreateModelResources(sc, std::move(external_file), tag_suffix); + } + + absl::StatusOr GetOrCreateModelResources( + SubgraphContext* sc, std::unique_ptr external_file, + 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 // model asset bundle resources into the model resources graph service on