Refactor Prepocessor/Postprocessor

Major change:
1. Created a new base class, Processor, to let Prepocessor/Postprocessor share common logic of creation, sanity check, metadata extracting etc.
2. Created a new factory method to build a subclass of Processor. Simplified Processor creation by encapsulating the constructor and SanityCheck together.
3. Updated method name to follow tflite_support conversion: Tensor() -> GetTensor(), Metadata() -> GetTensorMetadata().

PiperOrigin-RevId: 408735018

Author: lu-wang-g

tensorflow_lite_support/cc/task/processor/audio_preprocessor.cc

@@ -74,11 +74,10 @@ tflite::support::StatusOr<const AudioProperties*> GetAudioPropertiesSafe(
 tflite::support::StatusOr<std::unique_ptr<AudioPreprocessor>>
 AudioPreprocessor::Create(tflite::task::core::TfLiteEngine* engine,
                           const std::initializer_list<int> input_indices) {
-  RETURN_IF_ERROR(Preprocessor::SanityCheck(/* num_expected_tensors = */ 1,
-                                            engine, input_indices,
-                                            /* requires_metadata = */ true));
-  auto processor =
-      ::absl::WrapUnique(new AudioPreprocessor(engine, input_indices));
+  ASSIGN_OR_RETURN(auto processor,
+                   Processor::Create<AudioPreprocessor>(
+                       /* num_expected_tensors = */ 1, engine, input_indices));
+
   RETURN_IF_ERROR(processor->Init());
   return processor;
 }
@@ -90,8 +89,9 @@ absl::Status AudioPreprocessor::Init() {
 }
 
 absl::Status AudioPreprocessor::SetAudioFormatFromMetadata() {
-  ASSIGN_OR_RETURN(const AudioProperties* props,
-                   GetAudioPropertiesSafe(Metadata(), input_indices_.at(0)));
+  ASSIGN_OR_RETURN(
+      const AudioProperties* props,
+      GetAudioPropertiesSafe(GetTensorMetadata(), tensor_indices_.at(0)));
   audio_format_.channels = props->channels();
   audio_format_.sample_rate = props->sample_rate();
   if (audio_format_.channels <= 0 || audio_format_.sample_rate <= 0) {
@@ -106,16 +106,16 @@ absl::Status AudioPreprocessor::SetAudioFormatFromMetadata() {
 
 absl::Status AudioPreprocessor::CheckAndSetInputs() {
   input_buffer_size_ = 1;
-  for (int i = 0; i < Tensor()->dims->size; i++) {
-    if (Tensor()->dims->data[i] < 1) {
+  for (int i = 0; i < GetTensor()->dims->size; i++) {
+    if (GetTensor()->dims->data[i] < 1) {
       return CreateStatusWithPayload(
           absl::StatusCode::kInvalidArgument,
           absl::StrFormat("Invalid size: %d for input tensor dimension: %d.",
-                          Tensor()->dims->data[i], i),
+                          GetTensor()->dims->data[i], i),
           tflite::support::TfLiteSupportStatus::
               kInvalidInputTensorDimensionsError);
     }
-    input_buffer_size_ *= Tensor()->dims->data[i];
+    input_buffer_size_ *= GetTensor()->dims->data[i];
   }
   // Check if the input buffer size is divisible by the required audio channels.
   // This needs to be done after loading metadata and input.
@@ -158,7 +158,7 @@ absl::Status AudioPreprocessor::Preprocess(
         tflite::support::TfLiteSupportStatus::kInvalidArgumentError);
   }
   return tflite::task::core::PopulateTensor(audio_buffer.GetFloatBuffer(),
-                                            input_buffer_size_, Tensor());
+                                            input_buffer_size_, GetTensor());
 }
 
 }  // namespace processor

tensorflow_lite_support/cc/task/processor/classification_postprocessor.cc

@@ -44,11 +44,9 @@ tflite::support::StatusOr<std::unique_ptr<ClassificationPostprocessor>>
 ClassificationPostprocessor::Create(
     core::TfLiteEngine* engine, const std::initializer_list<int> output_indices,
     std::unique_ptr<ClassificationOptions> options) {
-  RETURN_IF_ERROR(Postprocessor::SanityCheck(/* num_expected_tensors = */ 1,
-                                             engine, output_indices));
-
-  auto processor =
-      absl::WrapUnique(new ClassificationPostprocessor(engine, output_indices));
+  ASSIGN_OR_RETURN(auto processor,
+                   Processor::Create<ClassificationPostprocessor>(
+                       /* num_expected_tensors = */ 1, engine, output_indices));
 
   RETURN_IF_ERROR(processor->Init(std::move(options)));
   return processor;
@@ -72,13 +70,13 @@ absl::Status ClassificationPostprocessor::Init(
         TfLiteSupportStatus::kInvalidArgumentError);
   }
 
-  ASSIGN_OR_RETURN(
-      classification_head_,
-      BuildClassificationHead(*engine_->metadata_extractor(), *Metadata(),
-                              options->display_names_locale()));
+  ASSIGN_OR_RETURN(classification_head_,
+                   BuildClassificationHead(*engine_->metadata_extractor(),
+                                           *GetTensorMetadata(),
+                                           options->display_names_locale()));
 
   // Sanity check output tensors
-  const TfLiteTensor* output_tensor = Tensor();
+  const TfLiteTensor* output_tensor = GetTensor();
   const int num_dimensions = output_tensor->dims->size;
   if (num_dimensions == 4) {
     if (output_tensor->dims->data[1] != 1 ||
@@ -87,7 +85,7 @@ absl::Status ClassificationPostprocessor::Init(
           StatusCode::kInvalidArgument,
           absl::StrFormat("Unexpected WxH sizes for output index %d: got "
                           "%dx%d, expected 1x1.",
-                          output_indices_.at(0), output_tensor->dims->data[2],
+                          tensor_indices_.at(0), output_tensor->dims->data[2],
                           output_tensor->dims->data[1]),
           TfLiteSupportStatus::kInvalidOutputTensorDimensionsError);
     }
@@ -98,15 +96,15 @@ absl::Status ClassificationPostprocessor::Init(
             "Unexpected number of dimensions for output index %d: got %dD, "
             "expected either 2D (BxN with B=1) or 4D (BxHxWxN with B=1, W=1, "
             "H=1).",
-            output_indices_.at(0), num_dimensions),
+            tensor_indices_.at(0), num_dimensions),
         TfLiteSupportStatus::kInvalidOutputTensorDimensionsError);
   }
   if (output_tensor->dims->data[0] != 1) {
     return CreateStatusWithPayload(
         StatusCode::kInvalidArgument,
         absl::StrFormat("The output array is expected to have a batch size "
                         "of 1. Got %d for output index %d.",
-                        output_tensor->dims->data[0], output_indices_.at(0)),
+                        output_tensor->dims->data[0], tensor_indices_.at(0)),
         TfLiteSupportStatus::kInvalidOutputTensorDimensionsError);
   }
   int num_classes = output_tensor->dims->data[num_dimensions - 1];
@@ -126,7 +124,7 @@ absl::Status ClassificationPostprocessor::Init(
         absl::StrFormat("Got %d class(es) for output index %d, expected %d "
                         "according to the label map.",
                         output_tensor->dims->data[num_dimensions - 1],
-                        output_indices_.at(0), num_label_map_items),
+                        tensor_indices_.at(0), num_label_map_items),
         TfLiteSupportStatus::kMetadataInconsistencyError);
   }
   if (output_tensor->type != kTfLiteUInt8 &&
@@ -156,7 +154,7 @@ absl::Status ClassificationPostprocessor::Init(
     if (head_class_names.empty()) {
       std::string name = classification_head_.name;
       if (name.empty()) {
-        name = absl::StrFormat("#%d", output_indices_.at(0));
+        name = absl::StrFormat("#%d", tensor_indices_.at(0));
       }
       return CreateStatusWithPayload(
           StatusCode::kInvalidArgument,

tensorflow_lite_support/cc/task/processor/classification_postprocessor.h

@@ -100,12 +100,12 @@ class ClassificationPostprocessor : public Postprocessor {
 template <typename T>
 absl::Status ClassificationPostprocessor::Postprocess(T* classifications) {
   const auto& head = classification_head_;
-  classifications->set_head_index(output_indices_.at(0));
+  classifications->set_head_index(tensor_indices_.at(0));
 
   std::vector<std::pair<int, float>> score_pairs;
   score_pairs.reserve(head.label_map_items.size());
 
-  const TfLiteTensor* output_tensor = Tensor();
+  const TfLiteTensor* output_tensor = GetTensor();
   if (output_tensor->type == kTfLiteUInt8) {
     ASSIGN_OR_RETURN(const uint8* output_data,
                      core::AssertAndReturnTypedTensor<uint8>(output_tensor));

tensorflow_lite_support/cc/task/processor/embedding_postprocessor.cc

@@ -24,12 +24,10 @@ tflite::support::StatusOr<std::unique_ptr<EmbeddingPostprocessor>>
 EmbeddingPostprocessor::Create(core::TfLiteEngine* engine,
                                const std::initializer_list<int> output_indices,
                                std::unique_ptr<EmbeddingOptions> options) {
-  RETURN_IF_ERROR(Postprocessor::SanityCheck(/* num_expected_tensors = */ 1,
-                                             engine, output_indices,
-                                             /* requires_metadata = */ false));
-
-  auto processor =
-      absl::WrapUnique(new EmbeddingPostprocessor(engine, output_indices));
+  ASSIGN_OR_RETURN(auto processor,
+                   Processor::Create<EmbeddingPostprocessor>(
+                       /* num_expected_tensors = */ 1, engine, output_indices,
+                       /* requires_metadata = */ false));
 
   RETURN_IF_ERROR(processor->Init(std::move(options)));
   return processor;
@@ -39,8 +37,8 @@ absl::Status EmbeddingPostprocessor::Init(
     std::unique_ptr<EmbeddingOptions> options) {
   options_ = std::move(options);
 
-  int output_index = output_indices_.at(0);
-  auto* output_tensor = Tensor();
+  int output_index = tensor_indices_.at(0);
+  auto* output_tensor = GetTensor();
   int num_dimensions = output_tensor->dims->size;
   if (num_dimensions == 4) {
     if (output_tensor->dims->data[1] != 1 ||

tensorflow_lite_support/cc/task/processor/embedding_postprocessor.h

@@ -78,15 +78,15 @@ class EmbeddingPostprocessor : public Postprocessor {
 
 template <typename T>
 absl::Status EmbeddingPostprocessor::Postprocess(T* embedding) {
-  embedding->set_output_index(output_indices_.at(0));
+  embedding->set_output_index(tensor_indices_.at(0));
   auto* feature_vector = embedding->mutable_feature_vector();
-  if (Tensor()->type == kTfLiteUInt8) {
+  if (GetTensor()->type == kTfLiteUInt8) {
     const uint8* output_data =
         engine_->interpreter()->typed_output_tensor<uint8>(
-            output_indices_.at(0));
+            tensor_indices_.at(0));
     // Get the zero_point and scale parameters from the tensor metadata.
     const int output_tensor_index =
-        engine_->interpreter()->outputs()[output_indices_.at(0)];
+        engine_->interpreter()->outputs()[tensor_indices_.at(0)];
     const TfLiteTensor* output_tensor =
         engine_->interpreter()->tensor(output_tensor_index);
     for (int j = 0; j < embedding_dimension_; ++j) {
@@ -98,7 +98,7 @@ absl::Status EmbeddingPostprocessor::Postprocess(T* embedding) {
     // Float
     const float* output_data =
         engine_->interpreter()->typed_output_tensor<float>(
-            output_indices_.at(0));
+            tensor_indices_.at(0));
     for (int j = 0; j < embedding_dimension_; ++j) {
       feature_vector->add_value_float(output_data[j]);
     }

tensorflow_lite_support/cc/task/processor/image_preprocessor.cc

@@ -38,12 +38,10 @@ tflite::support::StatusOr<std::unique_ptr<ImagePreprocessor>>
 ImagePreprocessor::Create(
     core::TfLiteEngine* engine, const std::initializer_list<int> input_indices,
     const vision::FrameBufferUtils::ProcessEngine& process_engine) {
-  RETURN_IF_ERROR(Preprocessor::SanityCheck(/* num_expected_tensors = */ 1,
-                                            engine, input_indices,
-                                            /* requires_metadata = */ false));
-
-  auto processor =
-      absl::WrapUnique(new ImagePreprocessor(engine, input_indices));
+  ASSIGN_OR_RETURN(auto processor,
+                   Processor::Create<ImagePreprocessor>(
+                       /* num_expected_tensors = */ 1, engine, input_indices,
+                       /* requires_metadata = */ false));
 
   RETURN_IF_ERROR(processor->Init(process_engine));
   return processor;
@@ -85,7 +83,7 @@ absl::Status ImagePreprocessor::Init(
   }
 
   // Determine if the input shape is resizable.
-  const TfLiteIntArray* dims_signature = Tensor()->dims_signature;
+  const TfLiteIntArray* dims_signature = GetTensor()->dims_signature;
 
   // Some fixed-shape models do not have dims_signature.
   if (dims_signature != nullptr && dims_signature->size > 2) {
@@ -152,26 +150,26 @@ absl::Status ImagePreprocessor::Preprocess(const FrameBuffer& frame_buffer,
   // If dynamic, it will re-dim the entire graph as per the input.
   if (is_height_mutable_ || is_width_mutable_) {
     engine_->interpreter()->ResizeInputTensorStrict(
-        0, {Tensor()->dims->data[0], input_specs_.image_height,
-            input_specs_.image_width, Tensor()->dims->data[3]});
+        0, {GetTensor()->dims->data[0], input_specs_.image_height,
+            input_specs_.image_width, GetTensor()->dims->data[3]});
 
     engine_->interpreter()->AllocateTensors();
   }
   // Then normalize pixel data (if needed) and populate the input tensor.
   switch (input_specs_.tensor_type) {
     case kTfLiteUInt8:
-      if (Tensor()->bytes != input_data_byte_size) {
+      if (GetTensor()->bytes != input_data_byte_size) {
         return tflite::support::CreateStatusWithPayload(
             absl::StatusCode::kInternal,
             "Size mismatch or unsupported padding bytes between pixel data "
             "and input tensor.");
       }
       // No normalization required: directly populate data.
       RETURN_IF_ERROR(tflite::task::core::PopulateTensor(
-          input_data, input_data_byte_size / sizeof(uint8), Tensor()));
+          input_data, input_data_byte_size / sizeof(uint8), GetTensor()));
       break;
     case kTfLiteFloat32: {
-      if (Tensor()->bytes / sizeof(float) !=
+      if (GetTensor()->bytes / sizeof(float) !=
           input_data_byte_size / sizeof(uint8)) {
         return tflite::support::CreateStatusWithPayload(
             absl::StatusCode::kInternal,
@@ -181,7 +179,7 @@ absl::Status ImagePreprocessor::Preprocess(const FrameBuffer& frame_buffer,
       // Normalize and populate.
       ASSIGN_OR_RETURN(
           float* normalized_input_data,
-          tflite::task::core::AssertAndReturnTypedTensor<float>(Tensor()));
+          tflite::task::core::AssertAndReturnTypedTensor<float>(GetTensor()));
       const tflite::task::vision::NormalizationOptions& normalization_options =
           input_specs_.normalization_options.value();
       for (int i = 0; i < normalization_options.num_values; i++) {

tensorflow_lite_support/cc/task/processor/processor.cc

@@ -18,72 +18,35 @@ namespace tflite {
 namespace task {
 namespace processor {
 
-/* static */
-absl::Status Preprocessor::SanityCheck(
-    int num_expected_tensors, core::TfLiteEngine* engine,
-    const std::initializer_list<int> input_indices, bool requires_metadata) {
-  if (input_indices.size() != num_expected_tensors) {
-    return support::CreateStatusWithPayload(
-        absl::StatusCode::kInvalidArgument,
-        absl::StrFormat("Preprocessor can handle %d tensors, "
-                        "got: %d tensors.",
-                        num_expected_tensors, input_indices.size()));
-  }
-  for (auto* p = input_indices.begin(); p < input_indices.end(); p++) {
-    int input_index = *p;
-    if (input_index < 0 ||
-        input_index >= engine->InputCount(engine->interpreter())) {
-      return support::CreateStatusWithPayload(
-          absl::StatusCode::kInvalidArgument,
-          absl::StrFormat(
-              "Invalid input_index: %d. Model has %d input tensors.",
-              input_index, engine->InputCount(engine->interpreter())));
-    }
-    if (requires_metadata) {
-      auto* metadata =
-          engine->metadata_extractor()->GetInputTensorMetadata(input_index);
-      if (metadata == nullptr) {
-        return CreateStatusWithPayload(
-            absl::StatusCode::kInvalidArgument,
-            absl::StrFormat("Input tensor %d is missing TensorMetadata.",
-                            input_index),
-            support::TfLiteSupportStatus::kMetadataNotFoundError);
-      }
-    }
-  }
+constexpr char Preprocessor::kInputTypeName[];
+constexpr char Postprocessor::kOutputTypeName[];
 
-  return absl::OkStatus();
-}
-
-/* static */
-absl::Status Postprocessor::SanityCheck(
-    int num_expected_tensors, core::TfLiteEngine* engine,
-    const std::initializer_list<int> output_indices, bool requires_metadata) {
-  if (output_indices.size() != num_expected_tensors) {
+absl::Status Processor::SanityCheck(int num_expected_tensors,
+                                    bool requires_metadata) {
+  const char* tensor_type_name = GetTensorTypeName();
+  if (tensor_indices_.size() != num_expected_tensors) {
     return support::CreateStatusWithPayload(
         absl::StatusCode::kInvalidArgument,
-        absl::StrFormat("Postprocessor can handle %d tensors, "
+        absl::StrFormat("Processor can handle %d tensors, "
                         "got: %d tensors.",
-                        num_expected_tensors, output_indices.size()));
+                        num_expected_tensors, tensor_indices_.size()));
   }
-  for (auto* p = output_indices.begin(); p < output_indices.end(); p++) {
-    int output_index = *p;
-    if (output_index < 0 ||
-        output_index >= engine->OutputCount(engine->interpreter())) {
+
+  int tensor_count = GetModelTensorCount();
+  for (int i = 0; i < tensor_indices_.size(); i++) {
+    int index = tensor_indices_.at(i);
+    if (index < 0 || index >= tensor_count) {
       return support::CreateStatusWithPayload(
           absl::StatusCode::kInvalidArgument,
-          absl::StrFormat(
-              "Invalid output_index: %d. Model has %d output tensors.",
-              output_index, engine->OutputCount(engine->interpreter())));
+          absl::StrFormat("Invalid tensor_index: %d. Model has %d %s tensors.",
+                          index, tensor_count, tensor_type_name));
     }
     if (requires_metadata) {
-      auto* metadata =
-          engine->metadata_extractor()->GetOutputTensorMetadata(output_index);
-      if (metadata == nullptr) {
+      if (GetTensorMetadata(i) == nullptr) {
         return CreateStatusWithPayload(
             absl::StatusCode::kInvalidArgument,
-            absl::StrFormat("Output tensor %d is missing TensorMetadata.",
-                            output_index),
+            absl::StrFormat("%s tensor %d is missing TensorMetadata.",
+                            tensor_type_name, index),
             support::TfLiteSupportStatus::kMetadataNotFoundError);
       }
     }