Image postprocessor

tensorflow_lite_support/cc/task/processor/BUILD

@@ -40,6 +40,24 @@ cc_library_with_tflite(
     ],
 )
 
+cc_library_with_tflite(
+    name = "image_postprocessor",
+    srcs = ["image_postprocessor.cc"],
+    hdrs = ["image_postprocessor.h"],
+    tflite_deps = [
+        ":processor",
+        "//tensorflow_lite_support/cc/task/vision/utils:image_tensor_specs",
+    ],
+    deps = [
+        "//tensorflow_lite_support/cc/port:status_macros",
+        "//tensorflow_lite_support/cc/port:statusor",
+        "//tensorflow_lite_support/cc/task/vision/core:frame_buffer",
+        "//tensorflow_lite_support/cc/task/vision/utils:frame_buffer_utils",
+        "//tensorflow_lite_support/cc/task/core:task_utils",
+        "@com_google_absl//absl/status",
+    ],
+)
+
 cc_library_with_tflite(
     name = "classification_postprocessor",
     srcs = ["classification_postprocessor.cc"],

tensorflow_lite_support/cc/task/processor/image_postprocessor.cc

@@ -0,0 +1,129 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_lite_support/cc/task/processor/image_postprocessor.h"
+
+namespace tflite {
+namespace task {
+namespace processor {
+
+namespace {
+
+using ::absl::StatusCode;
+using ::tflite::metadata::ModelMetadataExtractor;
+using ::tflite::support::CreateStatusWithPayload;
+using ::tflite::support::StatusOr;
+using ::tflite::support::TfLiteSupportStatus;
+
+}  // namespace
+
+/* static */
+tflite::support::StatusOr<std::unique_ptr<ImagePostprocessor>>
+ImagePostprocessor::Create(core::TfLiteEngine* engine, const int output_index,
+                           const int input_index) {
+  ASSIGN_OR_RETURN(auto processor,
+                   Processor::Create<ImagePostprocessor>(
+                       /* num_expected_tensors = */ 1, engine, {output_index},
+                       /* requires_metadata = */ false));
+
+  RETURN_IF_ERROR(processor->Init(input_index));
+  return processor;
+}
+
+absl::Status ImagePostprocessor::Init(const int input_index) {
+  if (input_index == -1) {
+    return tflite::support::CreateStatusWithPayload(
+        absl::StatusCode::kInvalidArgument,
+        absl::StrFormat("Input image tensor not set. Input index found: %d",
+                        input_index),
+        tflite::support::TfLiteSupportStatus::kInputTensorNotFoundError);
+  }
+  ASSIGN_OR_RETURN(
+      auto output_specs,
+      vision::BuildImageTensorSpecs(*engine_->interpreter(),
+                                    *engine_->metadata_extractor(), false));
+  options_ = std::make_unique<vision::NormalizationOptions>(
+      output_specs.normalization_options.value());
+  return absl::OkStatus();
+}
+
+absl::StatusOr<vision::FrameBuffer> ImagePostprocessor::Postprocess() {
+  has_uint8_output_ = GetTensor()->type == kTfLiteUInt8;
+  const int kRgbPixelBytes = 3;
+
+  vision::FrameBuffer::Dimension to_buffer_dimension = {
+      GetTensor()->dims->data[2], GetTensor()->dims->data[1]};
+  size_t output_byte_size =
+      GetBufferByteSize(to_buffer_dimension, vision::FrameBuffer::Format::kRGB);
+  std::vector<uint8> postprocessed_data(output_byte_size / sizeof(uint8), 0);
+
+  if (has_uint8_output_) {  // No denormalization required.
+    if (GetTensor()->bytes != output_byte_size) {
+      return tflite::support::CreateStatusWithPayload(
+          absl::StatusCode::kInternal,
+          "Size mismatch or unsupported padding bytes between pixel data "
+          "and output tensor.");
+    }
+    const uint8* output_data =
+        core::AssertAndReturnTypedTensor<uint8>(GetTensor()).value();
+    postprocessed_data.insert(postprocessed_data.begin(), &output_data[0],
+                              &output_data[output_byte_size / sizeof(uint8)]);
+  } else {  // Denormalize to [0, 255] range.
+    if (GetTensor()->bytes / sizeof(float) !=
+        output_byte_size / sizeof(uint8)) {
+      return tflite::support::CreateStatusWithPayload(
+          absl::StatusCode::kInternal,
+          "Size mismatch or unsupported padding bytes between pixel data "
+          "and output tensor.");
+    }
+
+    uint8* denormalized_output_data = postprocessed_data.data();
+    const float* output_data =
+        core::AssertAndReturnTypedTensor<float>(GetTensor()).value();
+    const auto norm_options = GetNormalizationOptions();
+
+    if (norm_options.num_values == 1) {
+      float mean_value = norm_options.mean_values[0];
+      float std_value = norm_options.std_values[0];
+
+      for (size_t i = 0; i < output_byte_size / sizeof(uint8);
+           ++i, ++denormalized_output_data, ++output_data) {
+        *denormalized_output_data = static_cast<uint8>(std::round(std::min(
+            255.f, std::max(0.f, (*output_data) * std_value + mean_value))));
+      }
+    } else {
+      for (size_t i = 0; i < output_byte_size / sizeof(uint8);
+           ++i, ++denormalized_output_data, ++output_data) {
+        *denormalized_output_data = static_cast<uint8>(std::round(std::min(
+            255.f,
+            std::max(0.f, (*output_data) * norm_options.std_values[i % 3] +
+                              norm_options.mean_values[i % 3]))));
+      }
+    }
+  }
+
+  vision::FrameBuffer::Plane postprocessed_plane = {
+      /*buffer=*/postprocessed_data.data(),
+      /*stride=*/{GetTensor()->dims->data[2] * kRgbPixelBytes, kRgbPixelBytes}};
+  auto postprocessed_frame_buffer =
+      vision::FrameBuffer::Create({postprocessed_plane}, to_buffer_dimension,
+                                  vision::FrameBuffer::Format::kRGB,
+                                  vision::FrameBuffer::Orientation::kTopLeft);
+  return *postprocessed_frame_buffer.get();
+}
+
+}  // namespace processor
+}  // namespace task
+}  // namespace tflite

tensorflow_lite_support/cc/task/processor/image_postprocessor.h

@@ -0,0 +1,70 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either exPostss or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_SUPPORT_CC_TASK_PROCESSOR_IMAGE_POSTPROCESSOR_H_
+#define TENSORFLOW_LITE_SUPPORT_CC_TASK_PROCESSOR_IMAGE_POSTPROCESSOR_H_
+
+#include "tensorflow_lite_support/cc/port/status_macros.h"
+#include "tensorflow_lite_support/cc/task/processor/processor.h"
+#include "tensorflow_lite_support/cc/task/vision/core/frame_buffer.h"
+#include "tensorflow_lite_support/cc/task/vision/utils/frame_buffer_utils.h"
+#include "tensorflow_lite_support/cc/task/core/task_utils.h"
+#include "tensorflow_lite_support/cc/task/vision/utils/image_tensor_specs.h"
+
+namespace tflite {
+namespace task {
+namespace processor {
+
+// Process the associated output image tensor and convert it to a FrameBuffer.
+// Requirement for the output tensor:
+//   (kTfLiteUInt8/kTfLiteFloat32)
+//    - image input of size `[batch x height x width x channels]`.
+//    - batch inference is not supported (`batch` is required to be 1).
+//    - only RGB inputs are supported (`channels` is required to be 3).
+//    - if type is kTfLiteFloat32, NormalizationOptions are required to be
+//      attached to the metadata for output de-normalization. Uses input metadata
+//      as fallback in case output metadata isn't provided.
+class ImagePostprocessor : public Postprocessor {
+ public:
+  static tflite::support::StatusOr<std::unique_ptr<ImagePostprocessor>>
+  Create(core::TfLiteEngine* engine,
+         const int output_index,
+         const int input_index = -1);
+
+  // Processes the output tensor to an RGB of FrameBuffer type.
+  // If output tensor is of type kTfLiteFloat32, denormalize it into [0 - 255]
+  // via normalization parameters.
+  absl::StatusOr<vision::FrameBuffer> Postprocess();
+
+ private:
+  using Postprocessor::Postprocessor;
+
+  // Whether the model features quantized inference type (QUANTIZED_UINT8). This
+  // is currently detected by checking if the output tensor data type is uint8.
+  bool has_uint8_output_;
+
+  std::unique_ptr<vision::NormalizationOptions> options_;
+
+  absl::Status Init(const int input_index);
+
+  const vision::NormalizationOptions& GetNormalizationOptions() {
+    return *options_.get();
+  }
+};
+}  // namespace processor
+}  // namespace task
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_SUPPORT_CC_TASK_PROCESSOR_IMAGE_POSTPROCESSOR_H_

tensorflow_lite_support/cc/task/processor/image_preprocessor.cc

@@ -72,9 +72,9 @@ absl::Status ImagePreprocessor::Init(
     const vision::FrameBufferUtils::ProcessEngine& process_engine) {
   frame_buffer_utils_ = vision::FrameBufferUtils::Create(process_engine);
 
-  ASSIGN_OR_RETURN(input_specs_, vision::BuildInputImageTensorSpecs(
+  ASSIGN_OR_RETURN(input_specs_, vision::BuildImageTensorSpecs(
                                      *engine_->interpreter(),
-                                     *engine_->metadata_extractor()));
+                                     *engine_->metadata_extractor(), true));
 
   if (input_specs_.color_space != tflite::ColorSpaceType_RGB) {
     return tflite::support::CreateStatusWithPayload(