New image test utilities and memory management fixes.

PiperOrigin-RevId: 559926378
This commit is contained in:
MediaPipe Team 2023-08-24 18:00:46 -07:00 committed by Copybara-Service
parent dd09c8d3f7
commit 6e6978cdbf
3 changed files with 27 additions and 5 deletions

View File

@ -399,11 +399,13 @@ cc_library(
"//mediapipe/framework:packet", "//mediapipe/framework:packet",
"//mediapipe/framework:timestamp", "//mediapipe/framework:timestamp",
"//mediapipe/framework/formats:image", "//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_format_cc_proto",
"//mediapipe/framework/formats:image_frame", "//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_frame_opencv", "//mediapipe/framework/formats:image_frame_opencv",
"//mediapipe/framework/port:opencv_core", "//mediapipe/framework/port:opencv_core",
"//mediapipe/framework/port:opencv_imgcodecs", "//mediapipe/framework/port:opencv_imgcodecs",
"//mediapipe/framework/port:opencv_imgproc", "//mediapipe/framework/port:opencv_imgproc",
"@com_google_absl//absl/log",
], ],
) )

View File

@ -1,7 +1,15 @@
#include "mediapipe/util/image_test_utils.h" #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.h"
#include "mediapipe/framework/formats/image_frame_opencv.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_core_inc.h"
#include "mediapipe/framework/port/opencv_imgcodecs_inc.h" #include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
#include "mediapipe/framework/port/opencv_imgproc_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) { Packet MakeImageFramePacket(cv::Mat input, int timestamp) {
ImageFrame input_image(GetImageFormat(input.channels()), input.cols, ImageFrame input_image(GetImageFormat(input.channels()), input.cols,
input.rows, input.step, input.data, [](uint8_t*) {}); input.rows, input.step, input.data,
return MakePacket<ImageFrame>(std::move(input_image)).At(Timestamp(0)); [input](uint8_t*) mutable { input.release(); });
return MakePacket<ImageFrame>(std::move(input_image))
.At(Timestamp(timestamp));
} }
Packet MakeImagePacket(cv::Mat input, int timestamp) { Packet MakeImagePacket(cv::Mat input, int timestamp) {
mediapipe::Image input_image(std::make_shared<mediapipe::ImageFrame>( mediapipe::Image input_image(std::make_shared<mediapipe::ImageFrame>(
GetImageFormat(input.channels()), input.cols, input.rows, input.step, GetImageFormat(input.channels()), input.cols, input.rows, input.step,
input.data, [](uint8_t*) {})); input.data, [input](uint8_t*) mutable { input.release(); }));
return MakePacket<mediapipe::Image>(std::move(input_image)).At(Timestamp(0)); 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 } // namespace mediapipe

View File

@ -3,7 +3,7 @@
#include <string> #include <string>
#include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/packet.h" #include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/opencv_core_inc.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. // Converts the cv::Mat into Image packet.
Packet MakeImagePacket(cv::Mat input, int timestamp = 0); Packet MakeImagePacket(cv::Mat input, int timestamp = 0);
// Converts RGBA Mat to BGR.
cv::Mat RgbaToBgr(cv::Mat rgba);
} // namespace mediapipe } // namespace mediapipe
#endif // MEDIAPIPE_UTIL_IMAGE_TEST_UTILS_H_ #endif // MEDIAPIPE_UTIL_IMAGE_TEST_UTILS_H_