From a9c7e22ca4525289b17ad7471e813594b7870c01 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Wed, 9 Aug 2023 08:41:04 -0700
Subject: [PATCH] apply affine transform before drawing, in order to keep
 constant line width regardless of face cropping.

PiperOrigin-RevId: 555173659
---
 mediapipe/util/pose_util.cc | 17 +++++++++++------
 mediapipe/util/pose_util.h  |  5 +++--
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/mediapipe/util/pose_util.cc b/mediapipe/util/pose_util.cc
index 4a6bb6cdb..c68256cf8 100644
--- a/mediapipe/util/pose_util.cc
+++ b/mediapipe/util/pose_util.cc
@@ -192,15 +192,20 @@ void DrawPose(const mediapipe::NormalizedLandmarkList& pose, bool flip_y,
   }
 }
 
-void DrawFace(const mediapipe::NormalizedLandmarkList& face, bool flip_y,
-              bool draw_nose, int color_style, bool reverse_color,
+void DrawFace(const mediapipe::NormalizedLandmarkList& face,
+              const std::pair<int, int>& image_size, const cv::Mat& affine,
+              bool flip_y, bool draw_nose, int color_style, bool reverse_color,
               int draw_line_width, cv::Mat* image) {
-  const int target_width = image->cols;
-  const int target_height = image->rows;
   std::vector<cv::Point> landmarks;
   for (const auto& lm : face.landmark()) {
-    landmarks.emplace_back(lm.x() * target_width,
-                           (flip_y ? 1.0f - lm.y() : lm.y()) * target_height);
+    float ori_x = lm.x() * image_size.first;
+    float ori_y = (flip_y ? 1.0f - lm.y() : lm.y()) * image_size.second;
+
+    landmarks.emplace_back(
+        affine.at<float>(0, 0) * ori_x + affine.at<float>(0, 1) * ori_y +
+            affine.at<float>(0, 2),
+        affine.at<float>(1, 0) * ori_x + affine.at<float>(1, 1) * ori_y +
+            affine.at<float>(1, 2));
   }
 
   cv::Scalar kFaceOvalColor;
diff --git a/mediapipe/util/pose_util.h b/mediapipe/util/pose_util.h
index da952422f..aeb2b9222 100644
--- a/mediapipe/util/pose_util.h
+++ b/mediapipe/util/pose_util.h
@@ -23,8 +23,9 @@ namespace mediapipe {
 void DrawPose(const mediapipe::NormalizedLandmarkList& pose, bool flip_y,
               cv::Mat* image);
 
-void DrawFace(const mediapipe::NormalizedLandmarkList& face, bool flip_y,
-              bool draw_nose, int color_style, bool reverse_color,
+void DrawFace(const mediapipe::NormalizedLandmarkList& face,
+              const std::pair<int, int>& image_size, const cv::Mat& affine,
+              bool flip_y, bool draw_nose, int color_style, bool reverse_color,
               int draw_line_width, cv::Mat* image);
 
 }  // namespace mediapipe