absl::Status ImagePreprocessor::Preprocess()

in tensorflow_lite_support/cc/task/processor/image_preprocessor.cc [104:226]
101 lines of code
16 McCabe index (conditional complexity)

absl::Status ImagePreprocessor::Preprocess(const FrameBuffer& frame_buffer,
                                           const BoundingBox& roi) {
  // Input data to be normalized (if needed) and used for inference. In most
  // cases, this is the result of image preprocessing. In case no image
  // preprocessing is needed (see below), this points to the input frame
  // buffer raw data.
  const uint8* input_data;
  size_t input_data_byte_size;

  // Optional buffers in case image preprocessing is needed.
  std::unique_ptr<FrameBuffer> preprocessed_frame_buffer;
  std::vector<uint8> preprocessed_data;

  if (IsImagePreprocessingNeeded(frame_buffer, roi)) {
    // Preprocess input image to fit model requirements.
    // For now RGB is the only color space supported, which is ensured by
    // `InitInternal`.
    input_specs_.image_width =
        is_width_mutable_ ? roi.width() : input_specs_.image_width;
    input_specs_.image_height =
        is_height_mutable_ ? roi.height() : input_specs_.image_height;

    FrameBuffer::Dimension to_buffer_dimension = {input_specs_.image_width,
                                                  input_specs_.image_height};
    input_data_byte_size =
        GetBufferByteSize(to_buffer_dimension, FrameBuffer::Format::kRGB);
    preprocessed_data.resize(input_data_byte_size / sizeof(uint8), 0);
    input_data = preprocessed_data.data();

    FrameBuffer::Plane preprocessed_plane = {
        /*buffer=*/preprocessed_data.data(),
        /*stride=*/{input_specs_.image_width * kRgbPixelBytes, kRgbPixelBytes}};
    preprocessed_frame_buffer = FrameBuffer::Create(
        {preprocessed_plane}, to_buffer_dimension, FrameBuffer::Format::kRGB,
        FrameBuffer::Orientation::kTopLeft);

    RETURN_IF_ERROR(frame_buffer_utils_->Preprocess(
        frame_buffer, roi, preprocessed_frame_buffer.get()));
  } else {
    // Input frame buffer already targets model requirements: skip image
    // preprocessing. For RGB, the data is always stored in a single plane.
    input_data = frame_buffer.plane(0).buffer;
    input_data_byte_size = frame_buffer.plane(0).stride.row_stride_bytes *
                           frame_buffer.dimension().height;
  }

  // If dynamic, it will re-dim the entire graph as per the input.
  if (is_height_mutable_ || is_width_mutable_) {
    engine_->interpreter()->ResizeInputTensorStrict(
        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 (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), GetTensor()));
      break;
    case kTfLiteFloat32: {
      if (GetTensor()->bytes / sizeof(float) !=
          input_data_byte_size / sizeof(uint8)) {
        return tflite::support::CreateStatusWithPayload(
            absl::StatusCode::kInternal,
            "Size mismatch or unsupported padding bytes between pixel data "
            "and input tensor.");
      }
      // Normalize and populate.
      ASSIGN_OR_RETURN(
          float* normalized_input_data,
          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++) {
        if (std::abs(normalization_options.std_values[i]) <
            std::numeric_limits<float>::epsilon()) {
          return tflite::support::CreateStatusWithPayload(
              absl::StatusCode::kInternal,
              "NormalizationOptions.std_values can't be 0. Please check if the "
              "tensor metadata has been populated correctly.");
        }
      }
      if (normalization_options.num_values == 1) {
        float mean_value = normalization_options.mean_values[0];
        float inv_std_value = (1.0f / normalization_options.std_values[0]);
        for (size_t i = 0; i < input_data_byte_size / sizeof(uint8);
             i++, input_data++, normalized_input_data++) {
          *normalized_input_data =
              inv_std_value * (static_cast<float>(*input_data) - mean_value);
        }
      } else {
        std::array<float, 3> inv_std_values = {
            1.0f / normalization_options.std_values[0],
            1.0f / normalization_options.std_values[1],
            1.0f / normalization_options.std_values[2]};
        for (size_t i = 0; i < input_data_byte_size / sizeof(uint8);
             i++, input_data++, normalized_input_data++) {
          *normalized_input_data = inv_std_values[i % 3] *
                                   (static_cast<float>(*input_data) -
                                    normalization_options.mean_values[i % 3]);
        }
      }
      break;
    }
    case kTfLiteInt8:
      return tflite::support::CreateStatusWithPayload(
          absl::StatusCode::kUnimplemented,
          "kTfLiteInt8 input type is not implemented yet.");
    default:
      return tflite::support::CreateStatusWithPayload(
          absl::StatusCode::kInternal, "Unexpected input tensor type.");
  }

  return absl::OkStatus();
}