From 9a8259b8cf1e4da72b06c49830171b7b3d684b7e Mon Sep 17 00:00:00 2001
From: Kinar <kinar.ravishankar@codeandtheory.com>
Date: Tue, 9 Jan 2024 13:42:02 -0800
Subject: [PATCH] Added code to free up the image buffer and included a test to
 verify it

---
 .../tasks/c/vision/pose_landmarker/BUILD      | 10 ++++++++
 .../pose_landmarker_result_converter.cc       |  8 ++++++-
 .../pose_landmarker_result_converter_test.cc  | 23 +++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/mediapipe/tasks/c/vision/pose_landmarker/BUILD b/mediapipe/tasks/c/vision/pose_landmarker/BUILD
index 2dbfba3070..a947e5a286 100644
--- a/mediapipe/tasks/c/vision/pose_landmarker/BUILD
+++ b/mediapipe/tasks/c/vision/pose_landmarker/BUILD
@@ -42,14 +42,24 @@ cc_library(
 cc_test(
     name = "pose_landmarker_result_converter_test",
     srcs = ["pose_landmarker_result_converter_test.cc"],
+    data = [
+        "//mediapipe/framework/formats:image_frame_opencv",
+        "//mediapipe/framework/port:opencv_core",
+        "//mediapipe/framework/port:opencv_imgproc",
+        "//mediapipe/tasks/testdata/vision:test_images",
+    ],
     linkstatic = 1,
     deps = [
         ":pose_landmarker_result",
         ":pose_landmarker_result_converter",
         "//mediapipe/framework/port:gtest",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/formats:image",
+        "//mediapipe/framework/deps:file_path",
         "//mediapipe/tasks/c/components/containers:landmark",
         "//mediapipe/tasks/cc/components/containers:landmark",
         "//mediapipe/tasks/cc/vision/pose_landmarker:pose_landmarker_result",
+        "//mediapipe/tasks/cc/vision/utils:image_utils",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter.cc b/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter.cc
index 88b40a37fc..267efdf358 100644
--- a/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter.cc
+++ b/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter.cc
@@ -78,9 +78,15 @@ void CppConvertToPoseLandmarkerResult(
 
 void CppClosePoseLandmarkerResult(PoseLandmarkerResult* result) {
   if (result->segmentation_masks) {
+    for (uint32_t i = 0; i < result->segmentation_masks_count; ++i) {
+      if (result->segmentation_masks[i].type == MpImage::IMAGE_FRAME) {
+        const uint8_t* buffer =
+            result->segmentation_masks[i].image_frame.image_buffer;
+      }
+      // TODO: Add similar cleanup for GPU_BUFFER later
+    }
     delete[] result->segmentation_masks;
     result->segmentation_masks = nullptr;
-    result->segmentation_masks_count = 0;
   }
 
   for (uint32_t i = 0; i < result->pose_landmarks_count; ++i) {
diff --git a/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter_test.cc b/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter_test.cc
index b1bb6cfc95..49ecb2d4f2 100644
--- a/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter_test.cc
+++ b/mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result_converter_test.cc
@@ -17,14 +17,24 @@ limitations under the License.
 
 #include <cstdint>
 
+#include "mediapipe/framework/deps/file_path.h"
+#include "mediapipe/framework/formats/image.h"
 #include "mediapipe/framework/port/gtest.h"
+#include "mediapipe/framework/port/status_matchers.h"
 #include "mediapipe/tasks/c/components/containers/landmark.h"
 #include "mediapipe/tasks/c/vision/pose_landmarker/pose_landmarker_result.h"
 #include "mediapipe/tasks/cc/components/containers/landmark.h"
 #include "mediapipe/tasks/cc/vision/pose_landmarker/pose_landmarker_result.h"
+#include "mediapipe/tasks/cc/vision/utils/image_utils.h"
 
 namespace mediapipe::tasks::c::components::containers {
 
+using ::mediapipe::file::JoinPath;
+using ::mediapipe::tasks::vision::DecodeImageFromFile;
+
+constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/";
+constexpr char kMaskImage[] = "segmentation_input_rotation0.jpg";
+
 void InitPoseLandmarkerResult(
     ::mediapipe::tasks::vision::pose_landmarker::PoseLandmarkerResult*
         cpp_result) {
@@ -46,6 +56,17 @@ void InitPoseLandmarkerResult(
   mediapipe::tasks::components::containers::Landmarks cpp_landmarks;
   cpp_landmarks.landmarks.push_back(cpp_landmark);
   cpp_result->pose_world_landmarks.push_back(cpp_landmarks);
+
+  // Initialize segmentation_masks
+  MP_ASSERT_OK_AND_ASSIGN(
+      Image mask_image,
+      DecodeImageFromFile(JoinPath("./", kTestDataDirectory, kMaskImage)));
+
+  // Ensure segmentation_masks is instantiated and add the mask_image to it.
+  if (!cpp_result->segmentation_masks) {
+    cpp_result->segmentation_masks.emplace();
+  }
+  cpp_result->segmentation_masks->push_back(mask_image);
 }
 
 TEST(PoseLandmarkerResultConverterTest, ConvertsCustomResult) {
@@ -102,11 +123,13 @@ TEST(PoseLandmarkerResultConverterTest, FreesMemory) {
 
   EXPECT_NE(c_result.pose_landmarks, nullptr);
   EXPECT_NE(c_result.pose_world_landmarks, nullptr);
+  EXPECT_NE(c_result.segmentation_masks, nullptr);
 
   CppClosePoseLandmarkerResult(&c_result);
 
   EXPECT_EQ(c_result.pose_landmarks, nullptr);
   EXPECT_EQ(c_result.pose_world_landmarks, nullptr);
+  EXPECT_EQ(c_result.segmentation_masks, nullptr);
 }
 
 }  // namespace mediapipe::tasks::c::components::containers