Added helper methods to retrieve input/output tensor types.

LiteRT-PiperOrigin-RevId: 760817334

Author: ai-edge-bot

litert/cc/litert_model.h

@@ -469,6 +469,48 @@ class Model : public internal::Handle<LiteRtModel, LiteRtDestroyModel> {
     return key;
   }
 
+  // Returns the tensor type for the given n-th input tensor.
+  Expected<RankedTensorType> GetInputTensorType(size_t signature_index,
+                                                size_t input_index) const {
+    auto subgraph = Subgraph(signature_index);
+    return subgraph->Inputs()[input_index].RankedTensorType();
+  }
+
+  // Returns the tensor type for the given input tensor name.
+  Expected<RankedTensorType> GetInputTensorType(
+      size_t signature_index, absl::string_view input_name) const {
+    auto subgraph = Subgraph(signature_index);
+    LITERT_ASSIGN_OR_RETURN(auto tensor, subgraph->Input(input_name));
+    return tensor.RankedTensorType();
+  }
+
+  // Get input tensor type of the default signature for input name.
+  Expected<RankedTensorType> GetInputTensorType(
+      absl::string_view input_name) const {
+    return GetInputTensorType(/*signature_index=*/0, input_name);
+  }
+
+  // Returns the tensor type for the given n-th output tensor.
+  Expected<RankedTensorType> GetOutputTensorType(size_t signature_index,
+                                                 size_t output_index) const {
+    auto subgraph = Subgraph(signature_index);
+    return subgraph->Outputs()[output_index].RankedTensorType();
+  }
+
+  // Returns the tensor type for the given output tensor name.
+  Expected<RankedTensorType> GetOutputTensorType(
+      size_t signature_index, absl::string_view output_name) const {
+    auto subgraph = Subgraph(signature_index);
+    LITERT_ASSIGN_OR_RETURN(auto tensor, subgraph->Output(output_name));
+    return tensor.RankedTensorType();
+  }
+
+  // Get output tensor type of the default signature for output name.
+  Expected<RankedTensorType> GetOutputTensorType(
+      absl::string_view output_name) const {
+    return GetOutputTensorType(/*signature_index=*/0, output_name);
+  }
+
  private:
   // Parameter `owned` indicates if the created TensorBuffer object should take
   // ownership of the provided `tensor_buffer` handle.

litert/samples/async_segmentation/segmentation_model.cc

@@ -48,19 +48,15 @@
 // buffer type.
 litert::Expected<std::vector<litert::TensorBuffer>> CreateGlInputBuffers(
     LiteRtEnvironment env, litert::CompiledModel& compiled_model,
-    litert::Signature& signature) {
-  LiteRtSubgraph subgraph_handle = signature.Subgraph();
-  litert::Subgraph subgraph = litert::Subgraph(subgraph_handle);
-
+    litert::Model& model, int signature_index) {
+  auto signature = model.GetSignature(signature_index);
   std::vector<litert::TensorBuffer> input_buffers;
-  input_buffers.reserve(subgraph.Inputs().size());
-  for (litert::Tensor& input_tensor : subgraph.Inputs()) {
+  for (int i = 0; i < signature->InputNames().size(); ++i) {
     LITERT_ASSIGN_OR_RETURN(
         litert::TensorBufferRequirements input_buffer_requirements,
-        compiled_model.GetInputBufferRequirements(signature.Key(),
-                                                  input_tensor.Name()));
+        compiled_model.GetInputBufferRequirements(signature_index, i));
     LITERT_ASSIGN_OR_RETURN(litert::RankedTensorType ranked_tensor_type,
-                            input_tensor.RankedTensorType());
+                            model.GetInputTensorType(signature_index, i));
     LITERT_ASSIGN_OR_RETURN(size_t buffer_size,
                             input_buffer_requirements.BufferSize());
     LITERT_ASSIGN_OR_RETURN(auto input_buffer,
@@ -76,21 +72,19 @@ litert::Expected<std::vector<litert::TensorBuffer>> CreateGlInputBuffers(
 // buffer type.
 litert::Expected<std::vector<litert::TensorBuffer>> CreateGlOutputBuffers(
     LiteRtEnvironment env, litert::CompiledModel& compiled_model,
-    litert::Signature& signature) {
-  LiteRtSubgraph subgraph_handle = signature.Subgraph();
-  litert::Subgraph subgraph = litert::Subgraph(subgraph_handle);
+    litert::Model& model, int signature_index) {
+  auto signature = model.GetSignature(signature_index);
 
   std::vector<litert::TensorBuffer> output_buffers;
-  output_buffers.reserve(subgraph.Outputs().size());
-  for (litert::Tensor& output_tensor : subgraph.Outputs()) {
+  output_buffers.reserve(signature->OutputNames().size());
+  for (int i = 0; i < signature->OutputNames().size(); ++i) {
     LITERT_ASSIGN_OR_RETURN(
-        litert::TensorBufferRequirements input_buffer_requirements,
-        compiled_model.GetOutputBufferRequirements(signature.Key(),
-                                                   output_tensor.Name()));
+        litert::TensorBufferRequirements output_buffer_requirements,
+        compiled_model.GetOutputBufferRequirements(signature_index, i));
     LITERT_ASSIGN_OR_RETURN(litert::RankedTensorType ranked_tensor_type,
-                            output_tensor.RankedTensorType());
+                            model.GetOutputTensorType(signature_index, i));
     LITERT_ASSIGN_OR_RETURN(size_t buffer_size,
-                            input_buffer_requirements.BufferSize());
+                            output_buffer_requirements.BufferSize());
     LITERT_ASSIGN_OR_RETURN(auto output_buffer,
                             litert::TensorBuffer::CreateManaged(
                                 env, kLiteRtTensorBufferTypeGlBuffer,
@@ -157,20 +151,18 @@ bool SegmentationModel::InitializeModel(const std::string& model_path,
     }
   }
 
-  env_ = std::make_unique<litert::Environment>(std::move(env));
-
   LITERT_ASSIGN_OR_ABORT(auto signatures, model_.GetSignatures());
 
   size_t signature_index = 0;
 
   if (use_gl_buffers_) {
-    LITERT_ASSIGN_OR_ABORT(input_buffers_,
-                           CreateGlInputBuffers(env.Get(), compiled_model_,
-                                                signatures[signature_index]));
+    LITERT_ASSIGN_OR_ABORT(
+        input_buffers_, CreateGlInputBuffers(env.Get(), compiled_model_, model_,
+                                             signature_index));
 
     LITERT_ASSIGN_OR_ABORT(output_buffers_,
                            CreateGlOutputBuffers(env.Get(), compiled_model_,
-                                                 signatures[signature_index]));
+                                                 model_, signature_index));
 
   } else {
     LITERT_ASSIGN_OR_ABORT(input_buffers_,
@@ -179,6 +171,9 @@ bool SegmentationModel::InitializeModel(const std::string& model_path,
     LITERT_ASSIGN_OR_ABORT(
         output_buffers_, compiled_model_.CreateOutputBuffers(signature_index));
   }
+
+  env_ = std::make_unique<litert::Environment>(std::move(env));
+
   std::cout << "SegmentationModel: Model initialized." << std::endl;
   std::cout << "SegmentationModel: warming up model..." << std::endl;
   auto start_time = absl::Now();