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

Update model_task_graph to support multiple local model resources.

Browse Source 
PiperOrigin-RevId: 482917453
This commit is contained in:
MediaPipe Team 2022-10-21 16:41:07 -07:00 committed by Copybara-Service
parent 4a6c23a76a
commit ea1d85d811
2 changed files with 32 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -156,21 +156,24 @@ absl::StatusOr<CalculatorGraphConfig> ModelTaskGraph::GetConfig(
}
absl::StatusOr<const ModelResources*> ModelTaskGraph::CreateModelResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file) {
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> 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<const ModelResources*> ModelTaskGraph::CreateModelResources(
absl::StatusOr<const ModelAssetBundleResources*>
ModelTaskGraph::CreateModelAssetBundleResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file) {
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> 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)));

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

@ -19,6 +19,7 @@ limitations under the License.
#include <memory>
#include <string>
#include <utility>
#include <vector>
#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<const ModelResources*> CreateModelResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file);
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> 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<const ModelAssetBundleResources*>
CreateModelAssetBundleResources(
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file);
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> 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<ModelResources> local_model_resources_;
std::vector<std::unique_ptr<ModelResources>> local_model_resources_;
std::unique_ptr<ModelAssetBundleResources>
std::vector<std::unique_ptr<ModelAssetBundleResources>>
local_model_asset_bundle_resources_;
};