// 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 "mediapipe/util/tflite/operations/max_unpooling.h" #include "tensorflow/lite/kernels/internal/common.h" #include "tensorflow/lite/kernels/internal/tensor.h" #include "tensorflow/lite/kernels/padding.h" namespace mediapipe { namespace tflite_operations { namespace { constexpr int kDataInputTensor = 0; constexpr int kIndicesTensor = 1; constexpr int kOutputTensor = 0; inline void MaxUnpooling(const ::tflite::PoolParams& params, const ::tflite::RuntimeShape& input_shape, const float* input_data, const float* indices_data, const ::tflite::RuntimeShape& output_shape, float* output_data) { TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4); const int batches = MatchingDim(input_shape, 0, output_shape, 0); const int depth = MatchingDim(input_shape, 3, output_shape, 3); const int input_height = input_shape.Dims(1); const int input_width = input_shape.Dims(2); const int stride_height = params.stride_height; const int stride_width = params.stride_width; std::memset(output_data, 0, output_shape.FlatSize() * sizeof(float)); for (int batch = 0; batch < batches; ++batch) { for (int in_y = 0; in_y < input_height; ++in_y) { for (int in_x = 0; in_x < input_width; ++in_x) { for (int channel = 0; channel < depth; ++channel) { const auto input_offset = Offset(input_shape, batch, in_y, in_x, channel); int idx = indices_data[input_offset]; const int max_x = idx % params.filter_width; const int max_y = idx / params.filter_width; const int out_x = in_x * stride_width - params.padding_values.width + max_x; const int out_y = in_y * stride_height - params.padding_values.height + max_y; output_data[Offset(output_shape, batch, out_y, out_x, channel)] = input_data[input_offset]; } } } } } TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->custom_initial_data); TfLitePaddingValues* data_padding = reinterpret_cast(node->user_data); TF_LITE_ENSURE_EQ(context, ::tflite::NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, ::tflite::NumOutputs(node), 1); TfLiteTensor* output = ::tflite::GetOutput(context, node, kOutputTensor); TF_LITE_ENSURE(context, output != nullptr); const TfLiteTensor* input = ::tflite::GetInput(context, node, kDataInputTensor); TF_LITE_ENSURE(context, input != nullptr); const TfLiteTensor* indices = ::tflite::GetInput(context, node, kIndicesTensor); TF_LITE_ENSURE(context, indices != nullptr); TF_LITE_ENSURE_EQ(context, ::tflite::NumDimensions(indices), 4); TF_LITE_ENSURE_EQ(context, ::tflite::NumDimensions(input), 4); TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32); TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); TF_LITE_ENSURE_EQ(context, indices->type, kTfLiteFloat32); int batches = input->dims->data[0]; int height = input->dims->data[1]; int width = input->dims->data[2]; int channels_out = input->dims->data[3]; int out_width = width * params->filter_width; int out_height = height * params->filter_height; data_padding->height = ::tflite::ComputePadding( params->stride_height, 1, out_height, params->filter_height, height); data_padding->width = ::tflite::ComputePadding( params->stride_width, 1, out_width, params->filter_width, width); TfLiteIntArray* output_size = TfLiteIntArrayCreate(4); output_size->data[0] = batches; output_size->data[1] = out_height; output_size->data[2] = out_width; output_size->data[3] = channels_out; return context->ResizeTensor(context, output, output_size); } TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->custom_initial_data); TfLitePaddingValues* data_padding = reinterpret_cast(node->user_data); TfLiteTensor* output = ::tflite::GetOutput(context, node, kOutputTensor); TF_LITE_ENSURE(context, output != nullptr); const TfLiteTensor* input = ::tflite::GetInput(context, node, kDataInputTensor); TF_LITE_ENSURE(context, input != nullptr); const TfLiteTensor* indices = ::tflite::GetInput(context, node, kIndicesTensor); TF_LITE_ENSURE(context, indices != nullptr); float activation_min, activation_max; ::tflite::CalculateActivationRange(params->activation, &activation_min, &activation_max); ::tflite::PoolParams op_params; op_params.stride_height = params->stride_height; op_params.stride_width = params->stride_width; op_params.filter_height = params->filter_height; op_params.filter_width = params->filter_width; op_params.padding_values.height = data_padding->height; op_params.padding_values.width = data_padding->width; op_params.float_activation_min = activation_min; op_params.float_activation_max = activation_max; MaxUnpooling(op_params, ::tflite::GetTensorShape(input), ::tflite::GetTensorData(input), ::tflite::GetTensorData(indices), ::tflite::GetTensorShape(output), ::tflite::GetTensorData(output)); return kTfLiteOk; } } // namespace TfLiteRegistration* RegisterMaxUnpooling2D() { static TfLiteRegistration reg = { [](TfLiteContext*, const char*, size_t) -> void* { return new TfLitePaddingValues(); }, [](TfLiteContext*, void* buffer) -> void { delete reinterpret_cast(buffer); }, Prepare, Eval}; return ® } } // namespace tflite_operations } // namespace mediapipe