Branch LandmarksToTransformMatrix operation to LandmarksToTransformMatrixV2.

PiperOrigin-RevId: 310610035
Change-Id: I5a284b3b539ddff85ba15c6c08967cab600d3dc8
This commit is contained in:
A. Unique TensorFlower 2020-05-08 12:46:14 -07:00 committed by TensorFlower Gardener
parent 0a3199d1bd
commit 696f2a8bd7

View File

@ -2485,6 +2485,37 @@ class Landmarks2TransformMatrixOperationParser : public TFLiteOperationParser {
}
};
class Landmarks2TransformMatrixV2OperationParser
: public TFLiteOperationParser {
public:
absl::Status IsSupported(const TfLiteContext* context,
const TfLiteNode* tflite_node,
const TfLiteRegistration* registration) final {
return CheckInputsOutputs(context, tflite_node, /*runtime_inputs=*/1,
/*outputs=*/1);
}
absl::Status Parse(const TfLiteNode* tflite_node,
const TfLiteRegistration* registration,
GraphFloat32* graph, ObjectReader* reader) final {
Node* node = graph->NewNode();
RETURN_IF_ERROR(reader->AddInput(node, 0)); // landmarks
RETURN_IF_ERROR(reader->AddOutputs(node)); // transform matrix
const std::string op_name = "landmarks_to_transform_matrix_v2";
node->operation.type = op_name;
BHWC output_shape;
RETURN_IF_ERROR(
ParseCustomAttributes(op_name, tflite_node->custom_initial_data,
tflite_node->custom_initial_data_size,
&(node->operation.attributes), &output_shape));
auto output_value = graph->FindOutputs(node->id)[0];
output_value->tensor.shape = output_shape;
return absl::OkStatus();
}
};
class AlignmentPointsToTransformMatrixOperationParser
: public TFLiteOperationParser {
public:
@ -2712,6 +2743,9 @@ std::unique_ptr<TFLiteOperationParser> NewOperationParser(
if (custom_name == "Landmarks2TransformMatrix") {
return std::make_unique<Landmarks2TransformMatrixOperationParser>();
}
if (custom_name == "Landmarks2TransformMatrixV2") {
return std::make_unique<Landmarks2TransformMatrixV2OperationParser>();
}
if (custom_name == "AlignmentPointsToTransformMatrix") {
return std::make_unique<
AlignmentPointsToTransformMatrixOperationParser>();