Project import generated by Copybara.
GitOrigin-RevId: 283c1a295de0a53e47d7a94996bda0c52dcfd677
This commit is contained in:
MediaPipe Team
@@ -0,0 +1,555 @@
|
||||
// Copyright 2021 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/landmarks_to_transform_matrix.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "tensorflow/lite/delegates/gpu/common/mediapipe/landmarks_to_transform_matrix.h"
|
||||
#include "tensorflow/lite/delegates/gpu/common/types.h"
|
||||
#include "tensorflow/lite/kernels/internal/common.h"
|
||||
#include "tensorflow/lite/kernels/internal/compatibility.h"
|
||||
#include "tensorflow/lite/kernels/internal/tensor.h"
|
||||
#include "tensorflow/lite/kernels/padding.h"
|
||||
#include "tensorflow/lite/schema/schema_generated.h"
|
||||
|
||||
using ::tflite::gpu::BHWC;
|
||||
using ::tflite::gpu::float2;
|
||||
using ::tflite::gpu::float3;
|
||||
using ::tflite::gpu::int2;
|
||||
using ::tflite::gpu::int3;
|
||||
using ::tflite::gpu::LandmarksToTransformMatrixV1Attributes;
|
||||
using ::tflite::gpu::LandmarksToTransformMatrixV2Attributes;
|
||||
|
||||
using ::tflite::GetInput;
|
||||
using ::tflite::GetOutput;
|
||||
using ::tflite::GetTensorData;
|
||||
using ::tflite::GetTensorShape;
|
||||
using ::tflite::NumDimensions;
|
||||
using ::tflite::NumInputs;
|
||||
using ::tflite::NumOutputs;
|
||||
using ::tflite::RuntimeShape;
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tflite_operations {
|
||||
namespace {
|
||||
|
||||
constexpr int kDataInputTensor = 0;
|
||||
constexpr int kOutputTensor = 0;
|
||||
constexpr int3 kTensformMatrixShape(1, 4, 4);
|
||||
|
||||
float2 Read3DLandmarkXY(const float* data, int idx) {
|
||||
float2 result;
|
||||
result.x = data[idx * 3];
|
||||
result.y = data[idx * 3 + 1];
|
||||
return result;
|
||||
}
|
||||
|
||||
float3 Read3DLandmarkXYZ(const float* data, int idx) {
|
||||
float3 result;
|
||||
result.x = data[idx * 3];
|
||||
result.y = data[idx * 3 + 1];
|
||||
result.z = data[idx * 3 + 2];
|
||||
return result;
|
||||
}
|
||||
|
||||
struct Mat3 {
|
||||
Mat3() { data.resize(9); }
|
||||
Mat3(float x00, float x01, float x02, float x10, float x11, float x12,
|
||||
float x20, float x21, float x22)
|
||||
: data{x00, x01, x02, x10, x11, x12, x20, x21, x22} {}
|
||||
|
||||
Mat3 operator*(const Mat3& other) {
|
||||
Mat3 result;
|
||||
for (int r = 0; r < 3; r++) {
|
||||
for (int c = 0; c < 3; c++) {
|
||||
float sum = 0;
|
||||
for (int k = 0; k < 3; k++) {
|
||||
sum += this->Get(r, k) * other.Get(k, c);
|
||||
}
|
||||
result.Set(r, c, sum);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
float3 operator*(const float3& vec) const {
|
||||
float3 result;
|
||||
for (int r = 0; r < 3; r++) {
|
||||
float sum = 0;
|
||||
for (int k = 0; k < 3; k++) {
|
||||
sum += this->Get(r, k) * vec[k];
|
||||
}
|
||||
result[r] = sum;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
float Get(int x, int y) const { return data[x * 3 + y]; }
|
||||
void Set(int x, int y, float val) { data[x * 3 + y] = val; }
|
||||
|
||||
std::vector<float> data;
|
||||
};
|
||||
|
||||
struct Mat4 {
|
||||
Mat4() { data.resize(16); }
|
||||
Mat4(float x00, float x01, float x02, float x03, float x10, float x11,
|
||||
float x12, float x13, float x20, float x21, float x22, float x23,
|
||||
float x30, float x31, float x32, float x33)
|
||||
: data{x00, x01, x02, x03, x10, x11, x12, x13,
|
||||
x20, x21, x22, x23, x30, x31, x32, x33} {}
|
||||
void operator*=(const Mat4& other) {
|
||||
Mat4 result;
|
||||
for (int r = 0; r < 4; r++) {
|
||||
for (int c = 0; c < 4; c++) {
|
||||
float sum = 0;
|
||||
for (int k = 0; k < 4; k++) {
|
||||
sum += this->Get(r, k) * other.Get(k, c);
|
||||
}
|
||||
result.Set(r, c, sum);
|
||||
}
|
||||
}
|
||||
std::memcpy(this->data.data(), result.data.data(),
|
||||
result.data.size() * sizeof(float));
|
||||
}
|
||||
float Get(int x, int y) const { return data[x * 4 + y]; }
|
||||
void Set(int x, int y, float val) { data[x * 4 + y] = val; }
|
||||
|
||||
std::vector<float> data;
|
||||
};
|
||||
|
||||
namespace v1 {
|
||||
|
||||
inline void LandmarksToTransformMatrixV1(
|
||||
const LandmarksToTransformMatrixV1Attributes& params,
|
||||
const RuntimeShape& input0_shape, const float* landmarks,
|
||||
const RuntimeShape& output_shape, float* output_data) {
|
||||
TFLITE_CHECK_EQ(input0_shape.DimensionsCount(), 4);
|
||||
TFLITE_CHECK_EQ(output_shape.DimensionsCount(), 3);
|
||||
TFLITE_CHECK_EQ(input0_shape.Dims(0), 1);
|
||||
TFLITE_CHECK_EQ(input0_shape.Dims(1), 1);
|
||||
TFLITE_CHECK_EQ(input0_shape.Dims(2), 1);
|
||||
|
||||
float2 left_landmark = Read3DLandmarkXY(landmarks, params.left_rotation_idx);
|
||||
float2 right_landmark =
|
||||
Read3DLandmarkXY(landmarks, params.right_rotation_idx);
|
||||
|
||||
float alpha = -std::atan((right_landmark.y - left_landmark.y) /
|
||||
(right_landmark.x - left_landmark.x));
|
||||
|
||||
float2 max_value(-100000, -100000);
|
||||
float2 min_value(100000, 100000);
|
||||
for (int i = 0; i < params.subset.size(); i++) {
|
||||
for (int j = 0; j < 2; j++) {
|
||||
float2 landmark_current =
|
||||
Read3DLandmarkXY(landmarks, params.subset[i][j]);
|
||||
float2 rotated(
|
||||
landmark_current.x * cos(alpha) - landmark_current.y * sin(alpha),
|
||||
landmark_current.x * sin(alpha) + landmark_current.y * cos(alpha));
|
||||
max_value = float2(std::max(max_value.x, rotated.x),
|
||||
std::max(max_value.y, rotated.y));
|
||||
min_value = float2(std::min(min_value.x, rotated.x),
|
||||
std::min(min_value.y, rotated.y));
|
||||
}
|
||||
}
|
||||
|
||||
float2 bbox_size((max_value.x - min_value.x) * params.bbox_size_multiplier,
|
||||
(max_value.y - min_value.y) * params.bbox_size_multiplier);
|
||||
|
||||
Mat3 scale_matrix(
|
||||
bbox_size.x / params.landmarks_range, 0.0, 0.0, // first row
|
||||
0.0, bbox_size.y / params.landmarks_range, 0.0, // second row
|
||||
0.0, 0.0, 1.0); // third row
|
||||
|
||||
float2 middle((max_value.x + min_value.x) / 2.0,
|
||||
(max_value.y + min_value.y) / 2.0);
|
||||
|
||||
float2 rotated_middle(middle.x * cos(-alpha) - middle.y * sin(-alpha),
|
||||
middle.x * sin(-alpha) + middle.y * cos(-alpha));
|
||||
|
||||
Mat3 rotation_matrix(
|
||||
cos(-alpha), -sin(-alpha),
|
||||
(rotated_middle.x / params.landmarks_range) * 2.0 - 1.0, // first row
|
||||
sin(-alpha), cos(-alpha),
|
||||
(rotated_middle.y / params.landmarks_range) * 2.0 - 1.0, // second row
|
||||
0, 0, 1); // third row
|
||||
|
||||
Mat3 to_relative(2.0 / (params.output_hw.w - 1.0), 0.0, -1.0, // first row
|
||||
0.0, 2.0 / (params.output_hw.h - 1.0), -1.0, // second row
|
||||
0.0, 0.0, 1.0); // third row
|
||||
|
||||
Mat3 to_absolute((params.input_hw.w - 1.0) / 2.0, 0.0,
|
||||
(params.input_hw.w - 1.0) / 2.0, // first row
|
||||
0.0, (params.input_hw.h - 1.0) / 2.0,
|
||||
(params.input_hw.h - 1.0) / 2.0, // second row
|
||||
0.0, 0.0, 1.0); // third row
|
||||
|
||||
// Inverse Transformstion Matrix
|
||||
Mat3 itm = to_absolute * rotation_matrix * scale_matrix * to_relative;
|
||||
|
||||
output_data[0] = itm.Get(0, 0);
|
||||
output_data[1] = itm.Get(0, 1);
|
||||
output_data[2] = 0.0;
|
||||
output_data[3] = itm.Get(0, 2);
|
||||
|
||||
output_data[4] = itm.Get(1, 0);
|
||||
output_data[5] = itm.Get(1, 1);
|
||||
output_data[6] = 0.0;
|
||||
output_data[7] = itm.Get(1, 2);
|
||||
|
||||
output_data[8] = itm.Get(2, 0);
|
||||
output_data[9] = itm.Get(2, 1);
|
||||
output_data[10] = itm.Get(2, 2);
|
||||
output_data[11] = 0.0;
|
||||
|
||||
output_data[12] = 0.0;
|
||||
output_data[13] = 0.0;
|
||||
output_data[14] = 0.0;
|
||||
output_data[15] = 1.0;
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
|
||||
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
||||
const TfLiteTensor* input = GetInput(context, node, kDataInputTensor);
|
||||
TF_LITE_ENSURE(context, input != nullptr);
|
||||
TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
|
||||
TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
|
||||
TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
|
||||
|
||||
TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
|
||||
output_size->data[0] = kTensformMatrixShape.x;
|
||||
output_size->data[1] = kTensformMatrixShape.y;
|
||||
output_size->data[2] = kTensformMatrixShape.z;
|
||||
|
||||
return context->ResizeTensor(context, output, output_size);
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
LandmarksToTransformMatrixV1Attributes op_params;
|
||||
BHWC output_shape;
|
||||
auto status = tflite::gpu::ParseLandmarksToTransformMatrixV1Attributes(
|
||||
node->custom_initial_data, node->custom_initial_data_size, &op_params,
|
||||
&output_shape);
|
||||
if (!status.ok()) {
|
||||
context->ReportError(context, status.message().data());
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.bbox_size_multiplier == 0) {
|
||||
context->ReportError(context, "Incorrect bbox_size_multiplier: %d",
|
||||
op_params.bbox_size_multiplier);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.dimensions != 3) {
|
||||
context->ReportError(context, "Incorrect dimensions: %d",
|
||||
op_params.dimensions);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.input_hw.h <= 0 || op_params.input_hw.w <= 0) {
|
||||
context->ReportError(context, "Incorrect input_hw: h = %d w = %d",
|
||||
op_params.input_hw.h, op_params.input_hw.w);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.output_hw.h <= 0 || op_params.output_hw.w <= 0) {
|
||||
context->ReportError(context, "Incorrect output_hw: h = %d w = %d",
|
||||
op_params.output_hw.h, op_params.output_hw.w);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.landmarks_range <= 0) {
|
||||
context->ReportError(context, "Incorrect landmarks_range: %d",
|
||||
op_params.landmarks_range);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.left_rotation_idx < 0) {
|
||||
context->ReportError(context, "Incorrect left_rotation_idx: %d",
|
||||
op_params.left_rotation_idx);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.right_rotation_idx < 0) {
|
||||
context->ReportError(context, "Incorrect right_rotation_idx: %d",
|
||||
op_params.right_rotation_idx);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.subset.empty()) {
|
||||
context->ReportError(context, "Subset parameter is empty");
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
int counter = 0;
|
||||
for (auto& val : op_params.subset) {
|
||||
for (int i = 0; i < 2; i++) {
|
||||
if (val[i] < 0) {
|
||||
context->ReportError(context,
|
||||
"Incorrect subset value: index = %d, value = %d",
|
||||
counter, val[i]);
|
||||
return kTfLiteError;
|
||||
}
|
||||
counter++;
|
||||
}
|
||||
}
|
||||
|
||||
const TfLiteTensor* input0 = GetInput(context, node, kDataInputTensor);
|
||||
TF_LITE_ENSURE(context, input0 != nullptr);
|
||||
TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
LandmarksToTransformMatrixV1(
|
||||
op_params, GetTensorShape(input0), GetTensorData<float>(input0),
|
||||
GetTensorShape(output), GetTensorData<float>(output));
|
||||
return kTfLiteOk;
|
||||
}
|
||||
} // namespace v1
|
||||
|
||||
namespace v2 {
|
||||
|
||||
void EstimateRotationRadians(const float* input_data_0, int left_rotation_idx,
|
||||
int right_rotation_idx,
|
||||
float target_rotation_radians,
|
||||
float* rotation_radians) {
|
||||
const float3 left_landmark =
|
||||
Read3DLandmarkXYZ(input_data_0, left_rotation_idx);
|
||||
const float3 right_landmark =
|
||||
Read3DLandmarkXYZ(input_data_0, right_rotation_idx);
|
||||
const float left_x = left_landmark[0];
|
||||
const float left_y = left_landmark[1];
|
||||
const float right_x = right_landmark[0];
|
||||
const float right_y = right_landmark[1];
|
||||
float rotation = std::atan2(right_y - left_y, right_x - left_x);
|
||||
rotation = target_rotation_radians - rotation;
|
||||
*rotation_radians = rotation;
|
||||
}
|
||||
|
||||
void EstimateCenterAndSize(const float* input_data_0,
|
||||
std::vector<tflite::gpu::int2> subset_idxs,
|
||||
float rotation_radians, float* crop_x, float* crop_y,
|
||||
float* crop_width, float* crop_height) {
|
||||
std::vector<float3> landmarks;
|
||||
landmarks.reserve(subset_idxs.size() * 2);
|
||||
for (int i = 0; i < subset_idxs.size(); i++) {
|
||||
landmarks.push_back(Read3DLandmarkXYZ(input_data_0, subset_idxs[i][0]));
|
||||
landmarks.push_back(Read3DLandmarkXYZ(input_data_0, subset_idxs[i][1]));
|
||||
}
|
||||
for (int i = 0; i < landmarks.size(); i++) {
|
||||
landmarks[i].z = 1.0;
|
||||
}
|
||||
const float& r = rotation_radians;
|
||||
// clang-format off
|
||||
const Mat3 t_rotation = Mat3(std::cos(r), -std::sin(r), 0.0,
|
||||
std::sin(r), std::cos(r), 0.0,
|
||||
0.0, 0.0, 1.0);
|
||||
const Mat3 t_rotation_inverse =
|
||||
Mat3(std::cos(-r), -std::sin(-r), 0.0,
|
||||
std::sin(-r), std::cos(-r), 0.0,
|
||||
0.0, 0.0, 1.0);
|
||||
// clang-format on
|
||||
for (int i = 0; i < landmarks.size(); i++) {
|
||||
landmarks[i] = t_rotation * landmarks[i];
|
||||
}
|
||||
float3 xy1_max = landmarks[0], xy1_min = landmarks[0];
|
||||
for (int i = 1; i < landmarks.size(); i++) {
|
||||
if (xy1_max.x < landmarks[i].x) xy1_max.x = landmarks[i].x;
|
||||
if (xy1_max.y < landmarks[i].y) xy1_max.y = landmarks[i].y;
|
||||
|
||||
if (xy1_min.x > landmarks[i].x) xy1_min.x = landmarks[i].x;
|
||||
if (xy1_min.y > landmarks[i].y) xy1_min.y = landmarks[i].y;
|
||||
}
|
||||
*crop_width = xy1_max.x - xy1_min.x;
|
||||
*crop_height = xy1_max.y - xy1_min.y;
|
||||
float3 crop_xy1 = xy1_min;
|
||||
crop_xy1.x += xy1_max.x;
|
||||
crop_xy1.y += xy1_max.y;
|
||||
crop_xy1.x /= 2;
|
||||
crop_xy1.y /= 2;
|
||||
crop_xy1 = t_rotation_inverse * crop_xy1;
|
||||
*crop_x = crop_xy1.x;
|
||||
*crop_y = crop_xy1.y;
|
||||
}
|
||||
|
||||
inline void LandmarksToTransformMatrixV2(
|
||||
const LandmarksToTransformMatrixV2Attributes& params,
|
||||
const RuntimeShape& input0_shape, const float* landmarks,
|
||||
const RuntimeShape& output_shape, float* output_data) {
|
||||
float rotation_radians = 0.0;
|
||||
EstimateRotationRadians(landmarks, params.left_rotation_idx,
|
||||
params.right_rotation_idx,
|
||||
params.target_rotation_radians, &rotation_radians);
|
||||
float crop_x = 0.0, crop_y = 0.0, crop_width = 0.0, crop_height = 0.0;
|
||||
EstimateCenterAndSize(landmarks, params.subset_idxs, rotation_radians,
|
||||
&crop_x, &crop_y, &crop_width, &crop_height);
|
||||
// Turn off clang formatting to make matrices initialization more readable.
|
||||
// clang-format off
|
||||
Mat4 t = Mat4(1.0, 0.0, 0.0, 0.0,
|
||||
0.0, 1.0, 0.0, 0.0,
|
||||
0.0, 0.0, 1.0, 0.0,
|
||||
0.0, 0.0, 0.0, 1.0);
|
||||
const Mat4 t_shift = Mat4(1.0, 0.0, 0.0, crop_x,
|
||||
0.0, 1.0, 0.0, crop_y,
|
||||
0.0, 0.0, 1.0, 0.0,
|
||||
0.0, 0.0, 0.0, 1.0);
|
||||
t *= t_shift;
|
||||
const float& r = -rotation_radians;
|
||||
const Mat4 t_rotation = Mat4(std::cos(r), -std::sin(r), 0.0, 0.0,
|
||||
std::sin(r), std::cos(r), 0.0, 0.0,
|
||||
0.0, 0.0, 1.0, 0.0,
|
||||
0.0, 0.0, 0.0, 1.0);
|
||||
t *= t_rotation;
|
||||
const float scale_x = params.scale_x * crop_width / params.output_width;
|
||||
const float scale_y = params.scale_y * crop_height / params.output_height;
|
||||
const Mat4 t_scale = Mat4(scale_x, 0.0, 0.0, 0.0,
|
||||
0.0, scale_y, 0.0, 0.0,
|
||||
0.0, 0.0, 1.0, 0.0,
|
||||
0.0, 0.0, 0.0, 1.0);
|
||||
t *= t_scale;
|
||||
const float shift_x = -1.0 * (params.output_width / 2.0);
|
||||
const float shift_y = -1.0 * (params.output_height / 2.0);
|
||||
const Mat4 t_shift2 = Mat4(1.0, 0.0, 0.0, shift_x,
|
||||
0.0, 1.0, 0.0, shift_y,
|
||||
0.0, 0.0, 1.0, 0.0,
|
||||
0.0, 0.0, 0.0, 1.0);
|
||||
t *= t_shift2;
|
||||
std::memcpy(output_data, t.data.data(), 16 * sizeof(float));
|
||||
// clang-format on
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
|
||||
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
||||
const TfLiteTensor* input = GetInput(context, node, kDataInputTensor);
|
||||
TF_LITE_ENSURE(context, input != nullptr);
|
||||
TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
TF_LITE_ENSURE_EQ(context, NumDimensions(input), 3);
|
||||
TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
|
||||
TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
|
||||
|
||||
TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
|
||||
output_size->data[0] = kTensformMatrixShape.x;
|
||||
output_size->data[1] = kTensformMatrixShape.y;
|
||||
output_size->data[2] = kTensformMatrixShape.z;
|
||||
|
||||
return context->ResizeTensor(context, output, output_size);
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
LandmarksToTransformMatrixV2Attributes op_params;
|
||||
BHWC output_shape;
|
||||
auto status = tflite::gpu::ParseLandmarksToTransformMatrixV2Attributes(
|
||||
node->custom_initial_data, node->custom_initial_data_size, &op_params,
|
||||
&output_shape);
|
||||
if (!status.ok()) {
|
||||
context->ReportError(context, status.message().data());
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.left_rotation_idx < 0) {
|
||||
context->ReportError(context, "Incorrect left_rotation_idx: %d",
|
||||
op_params.left_rotation_idx);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.right_rotation_idx < 0) {
|
||||
context->ReportError(context, "Incorrect right_rotation_idx: %d",
|
||||
op_params.right_rotation_idx);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.output_height <= 0) {
|
||||
context->ReportError(context, "Incorrect output_height: %d",
|
||||
op_params.output_height);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.output_width <= 0) {
|
||||
context->ReportError(context, "Incorrect output_width: %d",
|
||||
op_params.output_width);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.scale_x <= 0) {
|
||||
context->ReportError(context, "Incorrect scale_x: %d", op_params.scale_x);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
if (op_params.scale_y <= 0) {
|
||||
context->ReportError(context, "Incorrect scale_y: %d", op_params.scale_y);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
int counter = 0;
|
||||
for (auto& val : op_params.subset_idxs) {
|
||||
for (int i = 0; i < 2; i++) {
|
||||
if (val[i] < 0) {
|
||||
context->ReportError(context,
|
||||
"Incorrect subset value: index = %d, value = %d",
|
||||
counter, val[i]);
|
||||
return kTfLiteError;
|
||||
}
|
||||
counter++;
|
||||
}
|
||||
}
|
||||
|
||||
const TfLiteTensor* input0 = GetInput(context, node, kDataInputTensor);
|
||||
TF_LITE_ENSURE(context, input0 != nullptr);
|
||||
TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
LandmarksToTransformMatrixV2(
|
||||
op_params, GetTensorShape(input0), GetTensorData<float>(input0),
|
||||
GetTensorShape(output), GetTensorData<float>(output));
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
} // namespace v2
|
||||
|
||||
} // namespace
|
||||
|
||||
TfLiteRegistration* RegisterLandmarksToTransformMatrixV1() {
|
||||
static TfLiteRegistration reg = {
|
||||
/*.init=*/nullptr,
|
||||
/*.free=*/nullptr,
|
||||
/*.prepare=*/v1::Prepare,
|
||||
/*.invoke=*/v1::Eval,
|
||||
/*.profiling_string=*/nullptr,
|
||||
/*.builtin_code=*/tflite::BuiltinOperator_CUSTOM,
|
||||
/*.custom_name=*/"Landmarks2TransformMatrix",
|
||||
/*.version=*/1,
|
||||
};
|
||||
return ®
|
||||
}
|
||||
TfLiteRegistration* RegisterLandmarksToTransformMatrixV2() {
|
||||
static TfLiteRegistration reg = {
|
||||
/*.init=*/nullptr,
|
||||
/*.free=*/nullptr,
|
||||
/*.prepare=*/v2::Prepare,
|
||||
/*.invoke=*/v2::Eval,
|
||||
/*.profiling_string=*/nullptr,
|
||||
/*.builtin_code=*/tflite::BuiltinOperator_CUSTOM,
|
||||
/*.custom_name=*/"Landmarks2TransformMatrix",
|
||||
/*.version=*/2,
|
||||
};
|
||||
return ®
|
||||
}
|
||||
|
||||
} // namespace tflite_operations
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright 2021 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.
|
||||
|
||||
#ifndef MEDIAPIPE_UTIL_TFLITE_OPERATIONS_LANDMARKS_TO_TRANSFORM_MATRIX_H_
|
||||
#define MEDIAPIPE_UTIL_TFLITE_OPERATIONS_LANDMARKS_TO_TRANSFORM_MATRIX_H_
|
||||
|
||||
#include "tensorflow/lite/kernels/kernel_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tflite_operations {
|
||||
|
||||
TfLiteRegistration* RegisterLandmarksToTransformMatrixV1();
|
||||
|
||||
TfLiteRegistration* RegisterLandmarksToTransformMatrixV2();
|
||||
|
||||
} // namespace tflite_operations
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_UTIL_TFLITE_OPERATIONS_LANDMARKS_TO_TRANSFORM_MATRIX_H_
|
||||
@@ -0,0 +1,298 @@
|
||||
// Copyright 2021 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/transform_landmarks.h"
|
||||
|
||||
#include "tensorflow/lite/delegates/gpu/common/mediapipe/transform_landmarks.h"
|
||||
#include "tensorflow/lite/delegates/gpu/common/types.h"
|
||||
#include "tensorflow/lite/kernels/internal/common.h"
|
||||
#include "tensorflow/lite/kernels/internal/compatibility.h"
|
||||
#include "tensorflow/lite/kernels/internal/tensor.h"
|
||||
#include "tensorflow/lite/kernels/padding.h"
|
||||
#include "tensorflow/lite/schema/schema_generated.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tflite_operations {
|
||||
namespace {
|
||||
|
||||
constexpr int kDataInput0Tensor = 0;
|
||||
constexpr int kDataInput1Tensor = 1;
|
||||
constexpr int kOutputTensor = 0;
|
||||
|
||||
float DotProduct(const tflite::gpu::float4& l, const tflite::gpu::float4& r) {
|
||||
return l.x * r.x + l.y * r.y + l.z * r.z + l.w * r.w;
|
||||
}
|
||||
|
||||
namespace v1 {
|
||||
|
||||
inline void TransformLandmarks(
|
||||
const tflite::gpu::TransformLandmarksAttributes& params,
|
||||
const tflite::RuntimeShape& input0_shape, const float* landmarks,
|
||||
const tflite::RuntimeShape& input1_shape, const float* transform_matrix,
|
||||
const tflite::RuntimeShape& output_shape, float* output_data) {
|
||||
TFLITE_CHECK_EQ(input0_shape.DimensionsCount(), 4);
|
||||
TFLITE_CHECK_EQ(output_shape.DimensionsCount(), 4);
|
||||
const int output_height = output_shape.Dims(1);
|
||||
const int output_width = output_shape.Dims(2);
|
||||
const int output_channels = output_shape.Dims(3);
|
||||
TFLITE_CHECK_EQ(input0_shape.Dims(3) % params.dimensions, 0);
|
||||
TFLITE_CHECK_NE(params.scale, 0);
|
||||
|
||||
tflite::RuntimeShape input_shape_with_batch{/*batch=*/1, input0_shape.Dims(1),
|
||||
input0_shape.Dims(2),
|
||||
input0_shape.Dims(3)};
|
||||
tflite::RuntimeShape output_shape_with_batch{
|
||||
/*batch=*/1, output_shape.Dims(1), output_shape.Dims(2),
|
||||
output_shape.Dims(3)};
|
||||
|
||||
// Read first two rows of transformation matrix
|
||||
tflite::gpu::float4 x_transform(transform_matrix[0], transform_matrix[1],
|
||||
transform_matrix[2],
|
||||
transform_matrix[3] * params.scale);
|
||||
tflite::gpu::float4 y_transform(transform_matrix[4], transform_matrix[5],
|
||||
transform_matrix[6],
|
||||
transform_matrix[7] * params.scale);
|
||||
|
||||
for (int out_y = 0; out_y < output_height; ++out_y) {
|
||||
for (int out_x = 0; out_x < output_width; ++out_x) {
|
||||
for (int landmark = 0; landmark < output_channels / params.dimensions;
|
||||
++landmark) {
|
||||
const int offset = Offset(output_shape_with_batch, 0, out_y, out_x,
|
||||
landmark * params.dimensions);
|
||||
|
||||
if (params.dimensions == 2) {
|
||||
tflite::gpu::float4 lv(landmarks[offset], landmarks[offset + 1],
|
||||
static_cast<float>(0.0),
|
||||
static_cast<float>(1.0));
|
||||
tflite::gpu::float2 transformed(DotProduct(x_transform, lv),
|
||||
DotProduct(y_transform, lv));
|
||||
output_data[offset] = transformed.x;
|
||||
output_data[offset + 1] = transformed.y;
|
||||
}
|
||||
if (params.dimensions == 3) {
|
||||
tflite::gpu::float4 lv(landmarks[offset], landmarks[offset + 1],
|
||||
static_cast<float>(0.0),
|
||||
static_cast<float>(1.0));
|
||||
tflite::gpu::float3 transformed(DotProduct(x_transform, lv),
|
||||
DotProduct(y_transform, lv), lv.z);
|
||||
output_data[offset] = transformed.x;
|
||||
output_data[offset + 1] = transformed.y;
|
||||
output_data[offset + 2] = landmarks[offset + 2];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2);
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
|
||||
const TfLiteTensor* input =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
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);
|
||||
|
||||
TfLiteIntArray* output_size = TfLiteIntArrayCreate(4);
|
||||
output_size->data[0] = input->dims->data[0];
|
||||
output_size->data[1] = input->dims->data[1];
|
||||
output_size->data[2] = input->dims->data[2];
|
||||
output_size->data[3] = input->dims->data[3];
|
||||
|
||||
return context->ResizeTensor(context, output, output_size);
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
tflite::gpu::TransformLandmarksAttributes op_params;
|
||||
tflite::gpu::BHWC output_shape;
|
||||
auto status = tflite::gpu::ParseTransformLandmarksV1Attributes(
|
||||
node->custom_initial_data, node->custom_initial_data_size, &op_params,
|
||||
&output_shape);
|
||||
if (!status.ok()) {
|
||||
context->ReportError(context, status.message().data());
|
||||
return kTfLiteError;
|
||||
}
|
||||
if (op_params.dimensions != 3 && op_params.dimensions != 2) {
|
||||
context->ReportError(context, "Incorrect dimensions size: %d",
|
||||
op_params.dimensions);
|
||||
return kTfLiteError;
|
||||
}
|
||||
if (op_params.scale == 0) {
|
||||
context->ReportError(context, "Incorrect scale value: %d", op_params.scale);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
const TfLiteTensor* input0 =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input0 != nullptr);
|
||||
const TfLiteTensor* input1 =
|
||||
tflite::GetInput(context, node, kDataInput1Tensor);
|
||||
TF_LITE_ENSURE(context, input1 != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
TransformLandmarks(
|
||||
op_params, tflite::GetTensorShape(input0),
|
||||
tflite::GetTensorData<float>(input0), tflite::GetTensorShape(input1),
|
||||
tflite::GetTensorData<float>(input1), tflite::GetTensorShape(output),
|
||||
tflite::GetTensorData<float>(output));
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
} // namespace v1
|
||||
|
||||
namespace v2 {
|
||||
|
||||
inline void TransformLandmarksV2(
|
||||
const tflite::gpu::TransformLandmarksAttributes& params,
|
||||
const tflite::RuntimeShape& input0_shape, const float* landmarks,
|
||||
const float* transform_matrix, // transformation matrix
|
||||
const tflite::RuntimeShape& output_shape, float* output_data) {
|
||||
TFLITE_CHECK_EQ(input0_shape.DimensionsCount(), 3);
|
||||
TFLITE_CHECK_EQ(output_shape.DimensionsCount(), 3);
|
||||
const int output_width = output_shape.Dims(1);
|
||||
TFLITE_CHECK_EQ(input0_shape.Dims(2) % params.dimensions, 0);
|
||||
|
||||
tflite::RuntimeShape input_shape_with_batch{/*batch=*/1, input0_shape.Dims(0),
|
||||
input0_shape.Dims(1),
|
||||
input0_shape.Dims(2)};
|
||||
tflite::RuntimeShape output_shape_with_batch{
|
||||
/*batch=*/1, output_shape.Dims(0), output_shape.Dims(1),
|
||||
output_shape.Dims(2)};
|
||||
|
||||
// Read first two rows of transformation matrix
|
||||
tflite::gpu::float4 x_transform(transform_matrix[0], transform_matrix[1],
|
||||
transform_matrix[2], transform_matrix[3]);
|
||||
tflite::gpu::float4 y_transform(transform_matrix[4], transform_matrix[5],
|
||||
transform_matrix[6], transform_matrix[7]);
|
||||
|
||||
for (int landmark = 0; landmark < output_width; ++landmark) {
|
||||
const int offset = Offset(input_shape_with_batch, 0, 0, landmark, 0);
|
||||
|
||||
if (params.dimensions == 2) {
|
||||
tflite::gpu::float4 lv(landmarks[offset], landmarks[offset + 1],
|
||||
static_cast<float>(0.0), static_cast<float>(1.0));
|
||||
tflite::gpu::float2 transformed(DotProduct(x_transform, lv),
|
||||
DotProduct(y_transform, lv));
|
||||
output_data[offset] = transformed.x;
|
||||
output_data[offset + 1] = transformed.y;
|
||||
}
|
||||
if (params.dimensions == 3) {
|
||||
tflite::gpu::float4 lv(landmarks[offset], landmarks[offset + 1],
|
||||
static_cast<float>(0.0), static_cast<float>(1.0));
|
||||
tflite::gpu::float3 transformed(DotProduct(x_transform, lv),
|
||||
DotProduct(y_transform, lv), lv.z);
|
||||
output_data[offset] = transformed.x;
|
||||
output_data[offset + 1] = transformed.y;
|
||||
output_data[offset + 2] = landmarks[offset + 2];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2);
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
|
||||
const TfLiteTensor* input =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumDimensions(input), 3);
|
||||
TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
|
||||
TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
|
||||
|
||||
TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
|
||||
output_size->data[0] = input->dims->data[0];
|
||||
output_size->data[1] = input->dims->data[1];
|
||||
output_size->data[2] = input->dims->data[2];
|
||||
|
||||
return context->ResizeTensor(context, output, output_size);
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
tflite::gpu::TransformLandmarksAttributes op_params;
|
||||
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
tflite::RuntimeShape runtime_output_shape = tflite::GetTensorShape(output);
|
||||
tflite::gpu::BHWC output_shape(1, runtime_output_shape.Dims(0),
|
||||
runtime_output_shape.Dims(1),
|
||||
runtime_output_shape.Dims(2));
|
||||
auto status = tflite::gpu::ParseTransformLandmarksV2Attributes(
|
||||
node->custom_initial_data, node->custom_initial_data_size, &op_params,
|
||||
&output_shape);
|
||||
if (!status.ok()) {
|
||||
context->ReportError(context, status.message().data());
|
||||
return kTfLiteError;
|
||||
}
|
||||
if (op_params.dimensions != 3 && op_params.dimensions != 2) {
|
||||
context->ReportError(context, "Incorrect dimensions size: %d",
|
||||
op_params.dimensions);
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
const TfLiteTensor* input0 =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input0 != nullptr);
|
||||
const TfLiteTensor* input1 =
|
||||
tflite::GetInput(context, node, kDataInput1Tensor);
|
||||
TF_LITE_ENSURE(context, input1 != nullptr);
|
||||
|
||||
TransformLandmarksV2(op_params, tflite::GetTensorShape(input0),
|
||||
tflite::GetTensorData<float>(input0),
|
||||
tflite::GetTensorData<float>(input1),
|
||||
tflite::GetTensorShape(output),
|
||||
tflite::GetTensorData<float>(output));
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
} // namespace v2
|
||||
|
||||
} // namespace
|
||||
|
||||
TfLiteRegistration* RegisterTransformLandmarksV1() {
|
||||
static TfLiteRegistration reg = {
|
||||
/*.init=*/nullptr,
|
||||
/*.free=*/nullptr,
|
||||
/*.prepare=*/v1::Prepare,
|
||||
/*.invoke=*/v1::Eval,
|
||||
/*.profiling_string=*/nullptr,
|
||||
/*.builtin_code=*/tflite::BuiltinOperator_CUSTOM,
|
||||
/*.custom_name=*/"TransformLandmarks",
|
||||
/*.version=*/1,
|
||||
};
|
||||
return ®
|
||||
}
|
||||
|
||||
TfLiteRegistration* RegisterTransformLandmarksV2() {
|
||||
static TfLiteRegistration reg = {
|
||||
/*.init=*/nullptr,
|
||||
/*.free=*/nullptr,
|
||||
/*.prepare=*/v2::Prepare,
|
||||
/*.invoke=*/v2::Eval,
|
||||
/*.profiling_string=*/nullptr,
|
||||
/*.builtin_code=*/tflite::BuiltinOperator_CUSTOM,
|
||||
/*.custom_name=*/"TransformLandmarks",
|
||||
/*.version=*/2,
|
||||
};
|
||||
return ®
|
||||
}
|
||||
|
||||
} // namespace tflite_operations
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright 2021 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.
|
||||
|
||||
#ifndef MEDIAPIPE_UTIL_TFLITE_OPERATIONS_TRANSFORM_LANDMARKS_H_
|
||||
#define MEDIAPIPE_UTIL_TFLITE_OPERATIONS_TRANSFORM_LANDMARKS_H_
|
||||
|
||||
#include "tensorflow/lite/kernels/kernel_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tflite_operations {
|
||||
|
||||
TfLiteRegistration* RegisterTransformLandmarksV1();
|
||||
|
||||
TfLiteRegistration* RegisterTransformLandmarksV2();
|
||||
|
||||
} // namespace tflite_operations
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_UTIL_TFLITE_OPERATIONS_TRANSFORM_LANDMARKS_H_
|
||||
@@ -0,0 +1,332 @@
|
||||
// Copyright 2021 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/transform_tensor_bilinear.h"
|
||||
|
||||
#include "tensorflow/lite/delegates/gpu/common/mediapipe/transform_tensor_bilinear.h"
|
||||
#include "tensorflow/lite/delegates/gpu/common/types.h"
|
||||
#include "tensorflow/lite/kernels/internal/common.h"
|
||||
#include "tensorflow/lite/kernels/internal/compatibility.h"
|
||||
#include "tensorflow/lite/kernels/internal/tensor.h"
|
||||
#include "tensorflow/lite/kernels/padding.h"
|
||||
#include "tensorflow/lite/schema/schema_generated.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tflite_operations {
|
||||
namespace {
|
||||
|
||||
constexpr int kDataInput0Tensor = 0;
|
||||
constexpr int kDataInput1Tensor = 1;
|
||||
constexpr int kOutputTensor = 0;
|
||||
|
||||
float DotProduct(const tflite::gpu::float4& l, const tflite::gpu::float4& r) {
|
||||
return l.x * r.x + l.y * r.y + l.z * r.z + l.w * r.w;
|
||||
}
|
||||
namespace v1 {
|
||||
|
||||
inline void TransformTensor(
|
||||
const tflite::gpu::TransformTensorBilinearAttributes& params,
|
||||
const tflite::RuntimeShape& input0_shape,
|
||||
const float* input_data_0, // data
|
||||
const tflite::RuntimeShape& input1_shape,
|
||||
const float* input_data_1, // transformation matrix
|
||||
const tflite::RuntimeShape& output_shape, float* output_data) {
|
||||
TFLITE_CHECK_EQ(input0_shape.DimensionsCount(), 4);
|
||||
TFLITE_CHECK_EQ(output_shape.DimensionsCount(), 4);
|
||||
const int output_height = output_shape.Dims(1);
|
||||
const int output_width = output_shape.Dims(2);
|
||||
const int output_channels = output_shape.Dims(3);
|
||||
|
||||
const int input_height = input0_shape.Dims(1);
|
||||
const int input_width = input0_shape.Dims(2);
|
||||
const int input_channels = input0_shape.Dims(3);
|
||||
|
||||
tflite::RuntimeShape input_shape_with_batch{/*batch=*/1, input_height,
|
||||
input_width, input_channels};
|
||||
tflite::RuntimeShape output_shape_with_batch{/*batch=*/1, output_height,
|
||||
output_width, output_channels};
|
||||
|
||||
// Read first two rows of transformation matrix
|
||||
tflite::gpu::float4 x_transform(input_data_1[0], input_data_1[1],
|
||||
input_data_1[2], input_data_1[3]);
|
||||
tflite::gpu::float4 y_transform(input_data_1[4], input_data_1[5],
|
||||
input_data_1[6], input_data_1[7]);
|
||||
|
||||
for (int out_y = 0; out_y < output_height; ++out_y) {
|
||||
for (int out_x = 0; out_x < output_width; ++out_x) {
|
||||
tflite::gpu::float4 coord(
|
||||
static_cast<float>(out_x), static_cast<float>(out_y),
|
||||
static_cast<float>(0.0), static_cast<float>(1.0));
|
||||
|
||||
// Transformed coordinates.
|
||||
tflite::gpu::float2 tc(DotProduct(x_transform, coord),
|
||||
DotProduct(y_transform, coord));
|
||||
|
||||
bool out_of_bound = tc.x < 0.0 || tc.x > input_width - 1 || tc.y < 0.0 ||
|
||||
tc.y > input_height - 1;
|
||||
|
||||
for (int out_z = 0; out_z < output_channels; ++out_z) {
|
||||
float result = 0;
|
||||
if (!out_of_bound) {
|
||||
// Corners position:
|
||||
// q_11 --- q_21
|
||||
// ---- ----
|
||||
// q_12 --- q_22
|
||||
|
||||
auto ReadValue = [&](int h, int w) -> float {
|
||||
return h < 0 || w < 0 || h >= input_height || w >= input_width
|
||||
? 0
|
||||
: input_data_0[Offset(input_shape_with_batch, 0, h, w,
|
||||
out_z)];
|
||||
};
|
||||
|
||||
float q_11 = ReadValue(floor(tc.y), floor(tc.x));
|
||||
float q_21 = ReadValue(floor(tc.y), floor(tc.x) + 1);
|
||||
float q_12 = ReadValue(floor(tc.y) + 1, floor(tc.x));
|
||||
float q_22 = ReadValue(floor(tc.y) + 1, floor(tc.x) + 1);
|
||||
|
||||
float right_contrib = tc.x - floor(tc.x);
|
||||
float lower_contrib = tc.y - floor(tc.y);
|
||||
|
||||
float upper = (1.0 - right_contrib) * q_11 + right_contrib * q_21;
|
||||
float lower = (1.0 - right_contrib) * q_12 + right_contrib * q_22;
|
||||
|
||||
result = lower_contrib * lower + (1.0 - lower_contrib) * upper;
|
||||
}
|
||||
|
||||
const int out_offset =
|
||||
Offset(output_shape_with_batch, 0, out_y, out_x, out_z);
|
||||
|
||||
output_data[out_offset] = result;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2);
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
|
||||
const TfLiteTensor* input =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
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);
|
||||
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
tflite::gpu::TransformTensorBilinearAttributes op_params;
|
||||
tflite::gpu::BHWC output_shape;
|
||||
auto status = tflite::gpu::ParseTransformTensorBilinearV1Attributes(
|
||||
node->custom_initial_data, node->custom_initial_data_size, &op_params,
|
||||
&output_shape);
|
||||
if (!status.ok()) {
|
||||
context->ReportError(context, status.message().data());
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
const TfLiteTensor* input0 =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input0 != nullptr);
|
||||
const TfLiteTensor* input1 =
|
||||
tflite::GetInput(context, node, kDataInput1Tensor);
|
||||
TF_LITE_ENSURE(context, input1 != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
TransformTensor(
|
||||
op_params, tflite::GetTensorShape(input0),
|
||||
tflite::GetTensorData<float>(input0), tflite::GetTensorShape(input1),
|
||||
tflite::GetTensorData<float>(input1), tflite::GetTensorShape(output),
|
||||
tflite::GetTensorData<float>(output));
|
||||
return kTfLiteOk;
|
||||
}
|
||||
} // namespace v1
|
||||
|
||||
namespace v2 {
|
||||
|
||||
inline void TransformTensorBilinearV2(
|
||||
const tflite::gpu::TransformTensorBilinearAttributes& params,
|
||||
const tflite::RuntimeShape& input0_shape,
|
||||
const float* input_data_0, // data
|
||||
const tflite::RuntimeShape& input1_shape,
|
||||
const float* input_data_1, // transformation matrix
|
||||
const tflite::RuntimeShape& output_shape, float* output_data) {
|
||||
TFLITE_CHECK_EQ(input0_shape.DimensionsCount(), 4);
|
||||
TFLITE_CHECK_EQ(output_shape.DimensionsCount(), 4);
|
||||
const int output_height = output_shape.Dims(1);
|
||||
const int output_width = output_shape.Dims(2);
|
||||
const int output_channels = output_shape.Dims(3);
|
||||
|
||||
const int input_height = input0_shape.Dims(1);
|
||||
const int input_width = input0_shape.Dims(2);
|
||||
const int input_channels = input0_shape.Dims(3);
|
||||
|
||||
tflite::RuntimeShape input_shape_with_batch{/*batch=*/1, input_height,
|
||||
input_width, input_channels};
|
||||
tflite::RuntimeShape output_shape_with_batch{/*batch=*/1, output_height,
|
||||
output_width, output_channels};
|
||||
|
||||
// Read first two rows of transformation matrix
|
||||
tflite::gpu::float4 x_transform(input_data_1[0], input_data_1[1],
|
||||
input_data_1[2], input_data_1[3]);
|
||||
tflite::gpu::float4 y_transform(input_data_1[4], input_data_1[5],
|
||||
input_data_1[6], input_data_1[7]);
|
||||
|
||||
// Align corners correction: T -> S * ( T * A ), where T is a
|
||||
// transformation matrix, and subtruction and addition matrices are:
|
||||
// S A
|
||||
// 1 0 0 -0.5 1 0 0 0.5
|
||||
// 0 1 0 -0.5 0 1 0 0.5
|
||||
// 0 0 1 0 0 0 1 0
|
||||
// 0 0 0 1 0 0 0 1
|
||||
// Transformation matrix column 3 and rows 3, 4 are identity, which makes
|
||||
// the final formula pretty simple and easy to get if doing a manual
|
||||
// multiuplication.
|
||||
x_transform[3] += x_transform[0] * 0.5 + x_transform[1] * 0.5 - 0.5;
|
||||
y_transform[3] += y_transform[0] * 0.5 + y_transform[1] * 0.5 - 0.5;
|
||||
|
||||
for (int out_y = 0; out_y < output_height; ++out_y) {
|
||||
for (int out_x = 0; out_x < output_width; ++out_x) {
|
||||
tflite::gpu::float4 coord(
|
||||
static_cast<float>(out_x), static_cast<float>(out_y),
|
||||
static_cast<float>(0.0), static_cast<float>(1.0));
|
||||
|
||||
// Transformed coordinates.
|
||||
tflite::gpu::float2 tc(DotProduct(x_transform, coord),
|
||||
DotProduct(y_transform, coord));
|
||||
|
||||
bool out_of_bound = tc.x < 0.0 || tc.x > input_width - 1 || tc.y < 0.0 ||
|
||||
tc.y > input_height - 1;
|
||||
|
||||
for (int out_z = 0; out_z < output_channels; ++out_z) {
|
||||
float result = 0;
|
||||
if (!out_of_bound) {
|
||||
// Corners position:
|
||||
// q_11 --- q_21
|
||||
// ---- ----
|
||||
// q_12 --- q_22
|
||||
|
||||
auto ReadValue = [&](int h, int w) -> float {
|
||||
return h < 0 || w < 0 || h >= input_height || w >= input_width
|
||||
? 0
|
||||
: input_data_0[Offset(input_shape_with_batch, 0, h, w,
|
||||
out_z)];
|
||||
};
|
||||
|
||||
float q_11 = ReadValue(floor(tc.y), floor(tc.x));
|
||||
float q_21 = ReadValue(floor(tc.y), floor(tc.x) + 1);
|
||||
float q_12 = ReadValue(floor(tc.y) + 1, floor(tc.x));
|
||||
float q_22 = ReadValue(floor(tc.y) + 1, floor(tc.x) + 1);
|
||||
|
||||
float right_contrib = tc.x - floor(tc.x);
|
||||
float lower_contrib = tc.y - floor(tc.y);
|
||||
|
||||
float upper = (1.0 - right_contrib) * q_11 + right_contrib * q_21;
|
||||
float lower = (1.0 - right_contrib) * q_12 + right_contrib * q_22;
|
||||
|
||||
result = lower_contrib * lower + (1.0 - lower_contrib) * upper;
|
||||
}
|
||||
|
||||
const int out_offset =
|
||||
Offset(output_shape_with_batch, 0, out_y, out_x, out_z);
|
||||
|
||||
output_data[out_offset] = result;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2);
|
||||
TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
|
||||
const TfLiteTensor* input =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
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);
|
||||
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
tflite::gpu::TransformTensorBilinearAttributes op_params;
|
||||
tflite::gpu::BHWC output_shape;
|
||||
auto status = tflite::gpu::ParseTransformTensorBilinearV2Attributes(
|
||||
node->custom_initial_data, node->custom_initial_data_size, &op_params,
|
||||
&output_shape);
|
||||
if (!status.ok()) {
|
||||
context->ReportError(context, status.message().data());
|
||||
return kTfLiteError;
|
||||
}
|
||||
|
||||
const TfLiteTensor* input0 =
|
||||
tflite::GetInput(context, node, kDataInput0Tensor);
|
||||
TF_LITE_ENSURE(context, input0 != nullptr);
|
||||
const TfLiteTensor* input1 =
|
||||
tflite::GetInput(context, node, kDataInput1Tensor);
|
||||
TF_LITE_ENSURE(context, input1 != nullptr);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputTensor);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
|
||||
TransformTensorBilinearV2(
|
||||
op_params, tflite::GetTensorShape(input0),
|
||||
tflite::GetTensorData<float>(input0), tflite::GetTensorShape(input1),
|
||||
tflite::GetTensorData<float>(input1), tflite::GetTensorShape(output),
|
||||
tflite::GetTensorData<float>(output));
|
||||
return kTfLiteOk;
|
||||
}
|
||||
} // namespace v2
|
||||
|
||||
} // namespace
|
||||
|
||||
TfLiteRegistration* RegisterTransformTensorBilinearV1() {
|
||||
static TfLiteRegistration reg = {
|
||||
/*.init=*/nullptr,
|
||||
/*.free=*/nullptr,
|
||||
/*.prepare=*/v1::Prepare,
|
||||
/*.invoke=*/v1::Eval,
|
||||
/*.profiling_string=*/nullptr,
|
||||
/*.builtin_code=*/tflite::BuiltinOperator_CUSTOM,
|
||||
/*.custom_name=*/"TransformTensor",
|
||||
/*.version=*/1,
|
||||
};
|
||||
return ®
|
||||
}
|
||||
|
||||
TfLiteRegistration* RegisterTransformTensorBilinearV2() {
|
||||
static TfLiteRegistration reg = {
|
||||
/*.init=*/nullptr,
|
||||
/*.free=*/nullptr,
|
||||
/*.prepare=*/v2::Prepare,
|
||||
/*.invoke=*/v2::Eval,
|
||||
/*.profiling_string=*/nullptr,
|
||||
/*.builtin_code=*/tflite::BuiltinOperator_CUSTOM,
|
||||
/*.custom_name=*/"TransformTensorBilinear",
|
||||
/*.version=*/2,
|
||||
};
|
||||
return ®
|
||||
}
|
||||
|
||||
} // namespace tflite_operations
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright 2021 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.
|
||||
|
||||
#ifndef MEDIAPIPE_UTIL_TFLITE_OPERATIONS_TRANSFORM_TENSOR_BILINEAR_H_
|
||||
#define MEDIAPIPE_UTIL_TFLITE_OPERATIONS_TRANSFORM_TENSOR_BILINEAR_H_
|
||||
|
||||
#include "tensorflow/lite/kernels/kernel_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tflite_operations {
|
||||
|
||||
TfLiteRegistration* RegisterTransformTensorBilinearV1();
|
||||
|
||||
TfLiteRegistration* RegisterTransformTensorBilinearV2();
|
||||
|
||||
} // namespace tflite_operations
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_UTIL_TFLITE_OPERATIONS_TRANSFORM_TENSOR_BILINEAR_H_
|
||||
Reference in New Issue
Block a user