From 6e6978cdbf8509dc1f73932c811f95bf4e129f02 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Thu, 24 Aug 2023 18:00:46 -0700
Subject: [PATCH] New image test utilities and memory management fixes.

PiperOrigin-RevId: 559926378
---
 mediapipe/util/BUILD               |  2 ++
 mediapipe/util/image_test_utils.cc | 25 +++++++++++++++++++++----
 mediapipe/util/image_test_utils.h  |  5 ++++-
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/mediapipe/util/BUILD b/mediapipe/util/BUILD
index 9cf6f023f..9c655952a 100644
--- a/mediapipe/util/BUILD
+++ b/mediapipe/util/BUILD
@@ -399,11 +399,13 @@ cc_library(
         "//mediapipe/framework:packet",
         "//mediapipe/framework:timestamp",
         "//mediapipe/framework/formats:image",
+        "//mediapipe/framework/formats:image_format_cc_proto",
         "//mediapipe/framework/formats:image_frame",
         "//mediapipe/framework/formats:image_frame_opencv",
         "//mediapipe/framework/port:opencv_core",
         "//mediapipe/framework/port:opencv_imgcodecs",
         "//mediapipe/framework/port:opencv_imgproc",
+        "@com_google_absl//absl/log",
     ],
 )
 
diff --git a/mediapipe/util/image_test_utils.cc b/mediapipe/util/image_test_utils.cc
index 77b755953..c2966c567 100644
--- a/mediapipe/util/image_test_utils.cc
+++ b/mediapipe/util/image_test_utils.cc
@@ -1,7 +1,15 @@
 #include "mediapipe/util/image_test_utils.h"
 
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "absl/log/log.h"
+#include "mediapipe/framework/formats/image.h"
+#include "mediapipe/framework/formats/image_format.pb.h"
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/framework/formats/image_frame_opencv.h"
+#include "mediapipe/framework/packet.h"
 #include "mediapipe/framework/port/opencv_core_inc.h"
 #include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
 #include "mediapipe/framework/port/opencv_imgproc_inc.h"
@@ -43,15 +51,24 @@ mediapipe::ImageFormat::Format GetImageFormat(int image_channels) {
 
 Packet MakeImageFramePacket(cv::Mat input, int timestamp) {
   ImageFrame input_image(GetImageFormat(input.channels()), input.cols,
-                         input.rows, input.step, input.data, [](uint8_t*) {});
-  return MakePacket<ImageFrame>(std::move(input_image)).At(Timestamp(0));
+                         input.rows, input.step, input.data,
+                         [input](uint8_t*) mutable { input.release(); });
+  return MakePacket<ImageFrame>(std::move(input_image))
+      .At(Timestamp(timestamp));
 }
 
 Packet MakeImagePacket(cv::Mat input, int timestamp) {
   mediapipe::Image input_image(std::make_shared<mediapipe::ImageFrame>(
       GetImageFormat(input.channels()), input.cols, input.rows, input.step,
-      input.data, [](uint8_t*) {}));
-  return MakePacket<mediapipe::Image>(std::move(input_image)).At(Timestamp(0));
+      input.data, [input](uint8_t*) mutable { input.release(); }));
+  return MakePacket<mediapipe::Image>(std::move(input_image))
+      .At(Timestamp(timestamp));
+}
+
+cv::Mat RgbaToBgr(cv::Mat rgba) {
+  cv::Mat bgra;
+  cv::cvtColor(rgba, bgra, cv::COLOR_RGBA2BGR);
+  return bgra;
 }
 
 }  // namespace mediapipe
diff --git a/mediapipe/util/image_test_utils.h b/mediapipe/util/image_test_utils.h
index 6df9644d2..15a21c5b1 100644
--- a/mediapipe/util/image_test_utils.h
+++ b/mediapipe/util/image_test_utils.h
@@ -3,7 +3,7 @@
 
 #include <string>
 
-#include "mediapipe/framework/formats/image.h"
+#include "mediapipe/framework/formats/image_format.pb.h"
 #include "mediapipe/framework/packet.h"
 #include "mediapipe/framework/port/opencv_core_inc.h"
 
@@ -27,6 +27,9 @@ Packet MakeImageFramePacket(cv::Mat input, int timestamp = 0);
 // Converts the cv::Mat into Image packet.
 Packet MakeImagePacket(cv::Mat input, int timestamp = 0);
 
+// Converts RGBA Mat to BGR.
+cv::Mat RgbaToBgr(cv::Mat rgba);
+
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_UTIL_IMAGE_TEST_UTILS_H_