// Copyright 2019 The MediaPipe Authors.
//
// 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 <memory>

#include "mediapipe/calculators/tensorflow/image_frame_to_tensor_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/port/proto_ns.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_macros.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"

namespace mediapipe {

namespace tf = tensorflow;

namespace {
// Convert the ImageFrame into Tensor with floating point value type.
// The value will be normalized based on mean and stddev.
std::unique_ptr<tf::Tensor> ImageFrameToNormalizedTensor(
    // const ImageFrame& image_frame, float mean, float stddev) {
    const ImageFrame& image_frame,
    const mediapipe::proto_ns::RepeatedField<float>& mean,
    const mediapipe::proto_ns::RepeatedField<float>& stddev) {
  const int cols = image_frame.Width();
  const int rows = image_frame.Height();
  const int channels = image_frame.NumberOfChannels();
  const uint8* pixel = image_frame.PixelData();
  const int width_padding = image_frame.WidthStep() - cols * channels;
  auto tensor = ::absl::make_unique<tf::Tensor>(
      tf::DT_FLOAT, tf::TensorShape({rows, cols, channels}));
  auto tensor_data = tensor->tensor<float, 3>();

  for (int row = 0; row < rows; ++row) {
    for (int col = 0; col < cols; ++col) {
      for (int channel = 0; channel < channels; ++channel) {
        float mean_value = 0;
        if (mean.size() > 1) {
          mean_value = mean[channel];
        } else if (!mean.empty()) {
          mean_value = mean[0];
        }
        float stddev_value = 1;
        if (stddev.size() > 1) {
          stddev_value = stddev[channel];
        } else if (!stddev.empty()) {
          stddev_value = stddev[0];
        }
        tensor_data(row, col, channel) =
            (pixel[channel] - mean_value) / stddev_value;
      }
      pixel += channels;
    }
    pixel += width_padding;
  }
  return tensor;
}

}  // namespace

// Converts ImageFrames to TensorFlow Tensors.
//
// The calculator expects one input (a packet containing an ImageFrame) and
// generates one output (a packet containing a tf::Tensor holding the same
// pixel data). The output tensor will be 3D with dimensions corresponding to
// height, width, and the number of channels (e.g. 3 for RGB or 1 for GRAY8).
//
// This calculator supports ImageFrame objects with any valid format (SRGB
// SRGBA, GRAY8, GRAY16, and VEC32F1). It will generate a Tensor using DT_UINT8
// for the first three types, DT_UINT16 for GRAY16, and DT_FLOAT for VEC32F1.
//
// The ImageFrame data can be packed or padded. The pixel data will be copied
// to the Tensor in row-major order.
//
// Example config:
//  node {
//    calculator: "ImageFrameToTensorCalculator"
//    input_stream: "scaled_frames"
//    output_stream: "video_tensors"
//  }
class ImageFrameToTensorCalculator : public CalculatorBase {
 public:
  static absl::Status GetContract(CalculatorContract* cc);

  absl::Status Open(CalculatorContext* cc) override;
  absl::Status Process(CalculatorContext* cc) override;

 private:
  ImageFrameToTensorCalculatorOptions options_;
};
REGISTER_CALCULATOR(ImageFrameToTensorCalculator);

absl::Status ImageFrameToTensorCalculator::GetContract(CalculatorContract* cc) {
  // Start with only one input packet.
  RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
      << "Only one input stream is supported.";
  cc->Inputs().Index(0).Set<ImageFrame>(
      // ImageFrame frame.
  );
  RET_CHECK_EQ(cc->Outputs().NumEntries(), 1)
      << "Only one output stream is supported.";
  cc->Outputs().Index(0).Set<tf::Tensor>(
      // Output TensorFlow Tensor.
  );
  return absl::OkStatus();
}

absl::Status ImageFrameToTensorCalculator::Open(CalculatorContext* cc) {
  options_ = cc->Options<ImageFrameToTensorCalculatorOptions>();
  // Inform the framework that we always output at the same timestamp
  // as we receive a packet at.
  cc->SetOffset(TimestampDiff(0));
  return absl::OkStatus();
}

absl::Status ImageFrameToTensorCalculator::Process(CalculatorContext* cc) {
  const Packet& input_item = cc->Inputs().Index(0).Value();
  RET_CHECK(!input_item.IsEmpty()) << "Input cannot be empty.";

  // Extract the ImageFrame and metadata from the input packet.
  const ImageFrame& video_frame = input_item.Get<ImageFrame>();
  const int bytes_per_pixel = video_frame.ByteDepth();

  std::unique_ptr<tf::Tensor> tensor;
  if (options_.has_data_type()) {
    RET_CHECK_EQ(bytes_per_pixel, 1) << "Unsupported image format ("
                                     << bytes_per_pixel << " bytes per pixel)";
    const tf::DataType data_type = options_.data_type();
    RET_CHECK_EQ(data_type, tf::DT_FLOAT)
        << "Unsupported data type " << data_type;
    RET_CHECK_GT(options_.stddev().size(), 0) << "You must set a stddev.";
    RET_CHECK_GT(options_.stddev()[0], 0.0f) << "The stddev cannot be zero.";
    if (options_.stddev().size() > 1) {
      RET_CHECK_EQ(options_.stddev().size(), video_frame.NumberOfChannels())
          << "If specifying multiple stddev normalization values, "
          << "the number must match the number of image channels.";
    }
    if (options_.mean().size() > 1) {
      RET_CHECK_EQ(options_.mean().size(), video_frame.NumberOfChannels())
          << "If specifying multiple mean normalization values, "
          << "the number must match the number of image channels.";
    }
    tensor = ImageFrameToNormalizedTensor(video_frame, options_.mean(),
                                          options_.stddev());
  } else {
    const int height = video_frame.Height();
    const int width = video_frame.Width();
    const int num_channels = video_frame.NumberOfChannels();
    const int num_components = width * height * num_channels;
    tf::TensorShape tensor_shape({height, width, num_channels});

    // Use uint8 uint16, or float as the TF type depending on bpp of ImageFrame.
    tf::DataType data_type;
    if (bytes_per_pixel == 1) {
      data_type = tf::DT_UINT8;
    } else if (bytes_per_pixel == 2) {
      data_type = tf::DT_UINT16;
    } else if (bytes_per_pixel == 4) {
      data_type = tf::DT_FLOAT;
    } else {
      return absl::InvalidArgumentError(absl::StrCat(
          "Unsupported image format (", bytes_per_pixel, " bytes per pixel)"));
    }

    // This failure should never trigger, but it protects the code against
    // internal TF changes.
    RET_CHECK(tf::DataTypeCanUseMemcpy(data_type))
        << "Tensor data type does not support memcpy (type=" << data_type
        << ")";

    // Create the output tensor.
    tensor = ::absl::make_unique<tf::Tensor>(data_type, tensor_shape);

    // Copy pixel data from the ImageFrame to the tensor.
    if (data_type == tf::DT_UINT8) {
      uint8* dst = tensor->flat<uint8>().data();
      video_frame.CopyToBuffer(dst, num_components);
    } else if (data_type == tf::DT_UINT16) {
      uint16* dst = tensor->flat<uint16>().data();
      video_frame.CopyToBuffer(dst, num_components);
    } else {
      float* dst = tensor->flat<float>().data();
      video_frame.CopyToBuffer(dst, num_components);
    }
  }

  cc->Outputs().Index(0).Add(tensor.release(), cc->InputTimestamp());
  return absl::OkStatus();
}

}  // namespace mediapipe
