From 169bdf15b4a2de486059359cf1ecc02e289a72b4 Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Thu, 25 May 2023 11:18:45 -0700
Subject: [PATCH] Make an option for adjusting the face alignment output image
 size, and add a "transformation_matrix" output stream of the face stylizer
 graph.

PiperOrigin-RevId: 535319310
---
 .../face_stylizer/face_stylizer_graph.cc      | 21 +++++++++++++------
 .../proto/face_stylizer_graph_options.proto   |  3 +++
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/mediapipe/tasks/cc/vision/face_stylizer/face_stylizer_graph.cc b/mediapipe/tasks/cc/vision/face_stylizer/face_stylizer_graph.cc
index d4057467b..7fab54834 100644
--- a/mediapipe/tasks/cc/vision/face_stylizer/face_stylizer_graph.cc
+++ b/mediapipe/tasks/cc/vision/face_stylizer/face_stylizer_graph.cc
@@ -71,23 +71,21 @@ constexpr char kFaceLandmarksDetectorTFLiteName[] =
     "face_landmarks_detector.tflite";
 constexpr char kFaceStylizerTFLiteName[] = "face_stylizer.tflite";
 constexpr char kImageTag[] = "IMAGE";
-constexpr char kImageCpuTag[] = "IMAGE_CPU";
-constexpr char kImageGpuTag[] = "IMAGE_GPU";
 constexpr char kImageSizeTag[] = "IMAGE_SIZE";
 constexpr char kMatrixTag[] = "MATRIX";
 constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
 constexpr char kNormRectTag[] = "NORM_RECT";
-constexpr char kOutputSizeTag[] = "OUTPUT_SIZE";
 constexpr char kSizeTag[] = "SIZE";
 constexpr char kStylizedImageTag[] = "STYLIZED_IMAGE";
 constexpr char kTensorsTag[] = "TENSORS";
-constexpr int kFaceAlignmentOutputSize = 256;
+constexpr char kTransformationMatrixTag[] = "TRANSFORMATION_MATRIX";
 
 // Struct holding the different output streams produced by the face stylizer
 // graph.
 struct FaceStylizerOutputStreams {
   std::optional<Source<Image>> stylized_image;
   std::optional<Source<Image>> face_alignment_image;
+  std::optional<Source<std::array<float, 16>>> transformation_matrix;
   Source<Image> original_image;
 };
 
@@ -202,6 +200,10 @@ void ConfigureTensorsToImageCalculator(
 //     The aligned face image that is fed to the face stylization model to
 //     perform stylization. Also useful for preparing face stylization training
 //     data.
+//   TRANSFORMATION_MATRIX - std::array<float,16>
+//     An std::array<float, 16> representing a 4x4 row-major-order matrix that
+//     maps a point on the input image to a point on the output image, and
+//     can be used to reverse the mapping by inverting the matrix.
 //   IMAGE - mediapipe::Image
 //     The input image that the face landmarker runs on and has the pixel data
 //     stored on the target storage (CPU vs GPU).
@@ -280,6 +282,8 @@ class FaceStylizerGraph : public core::ModelTaskGraph {
       output_streams.face_alignment_image.value() >>
           graph[Output<Image>(kFaceAlignmentTag)];
     }
+    output_streams.transformation_matrix.value() >>
+        graph[Output<std::array<float, 16>>(kTransformationMatrixTag)];
     output_streams.original_image >> graph[Output<Image>(kImageTag)];
     return graph.GetConfig();
   }
@@ -357,9 +361,10 @@ class FaceStylizerGraph : public core::ModelTaskGraph {
           image_to_tensor.GetOptions<ImageToTensorCalculatorOptions>();
       image_to_tensor_options.mutable_output_tensor_float_range()->set_min(0);
       image_to_tensor_options.mutable_output_tensor_float_range()->set_max(1);
-      image_to_tensor_options.set_output_tensor_width(kFaceAlignmentOutputSize);
+      image_to_tensor_options.set_output_tensor_width(
+          task_options.face_alignment_size());
       image_to_tensor_options.set_output_tensor_height(
-          kFaceAlignmentOutputSize);
+          task_options.face_alignment_size());
       image_to_tensor_options.set_keep_aspect_ratio(true);
       image_to_tensor_options.set_border_mode(
           mediapipe::ImageToTensorCalculatorOptions::BORDER_ZERO);
@@ -378,6 +383,8 @@ class FaceStylizerGraph : public core::ModelTaskGraph {
 
       return {{/*stylized_image=*/std::nullopt,
                /*alignment_image=*/face_alignment,
+               /*transformation_matrix=*/
+               image_to_tensor.Out(kMatrixTag).Cast<std::array<float, 16>>(),
                /*original_image=*/pass_through.Out("").Cast<Image>()}};
     }
 
@@ -439,6 +446,8 @@ class FaceStylizerGraph : public core::ModelTaskGraph {
 
     return {{/*stylized_image=*/stylized,
              /*alignment_image=*/face_alignment,
+             /*transformation_matrix=*/
+             preprocessing.Out(kMatrixTag).Cast<std::array<float, 16>>(),
              /*original_image=*/preprocessing.Out(kImageTag).Cast<Image>()}};
   }
 };
diff --git a/mediapipe/tasks/cc/vision/face_stylizer/proto/face_stylizer_graph_options.proto b/mediapipe/tasks/cc/vision/face_stylizer/proto/face_stylizer_graph_options.proto
index 68abee43c..9528fab09 100644
--- a/mediapipe/tasks/cc/vision/face_stylizer/proto/face_stylizer_graph_options.proto
+++ b/mediapipe/tasks/cc/vision/face_stylizer/proto/face_stylizer_graph_options.proto
@@ -36,4 +36,7 @@ message FaceStylizerGraphOptions {
   // Options for face landmarker graph.
   optional vision.face_landmarker.proto.FaceLandmarkerGraphOptions
       face_landmarker_graph_options = 2;
+
+  // The width and height of the output face alignment images.
+  optional int32 face_alignment_size = 3 [default = 256];
 }