cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add an API in model_task_graph to create or use cached model resources.

Browse Source 
PiperOrigin-RevId: 559174528
This commit is contained in:
MediaPipe Team 2023-08-22 11:32:30 -07:00 committed by Copybara-Service
parent bcb83302bf
commit 8c4b971c14
2 changed files with 29 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
mediapipe/tasks/cc/core/model_task_graph.cc
View File

@ -186,6 +186,21 @@ absl::StatusOr<const ModelResources*> ModelTaskGraph::CreateModelResources(
return model_resources_cache_service.GetObject().GetModelResources(tag);
}
absl::StatusOr<const ModelResources*> ModelTaskGraph::GetOrCreateModelResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> 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<const ModelAssetBundleResources*>
ModelTaskGraph::CreateModelAssetBundleResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,

14
mediapipe/tasks/cc/core/model_task_graph.h
View File

@ -87,6 +87,20 @@ class ModelTaskGraph : public Subgraph {
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
std::string tag_suffix = "");
template <typename Options>
absl::StatusOr<const ModelResources*> GetOrCreateModelResources(
SubgraphContext* sc, std::string tag_suffix = "") {
auto external_file = std::make_unique<proto::ExternalFile>();
external_file->Swap(sc->MutableOptions<Options>()
->mutable_base_options()
->mutable_model_asset());
return GetOrCreateModelResources(sc, std::move(external_file), tag_suffix);
}
absl::StatusOr<const ModelResources*> GetOrCreateModelResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> 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