From 7a87546c30c347f8fc8d046431dbb27208a0f920 Mon Sep 17 00:00:00 2001
From: Camillo Lugaresi <camillol@google.com>
Date: Tue, 15 Nov 2022 14:35:54 -0800
Subject: [PATCH] Internal change

PiperOrigin-RevId: 488761646
---
 mediapipe/framework/tool/test_util.cc | 22 +++++++++++++---------
 mediapipe/framework/tool/test_util.h  |  4 ++++
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/mediapipe/framework/tool/test_util.cc b/mediapipe/framework/tool/test_util.cc
index 6433c93d2..c7ed063e0 100644
--- a/mediapipe/framework/tool/test_util.cc
+++ b/mediapipe/framework/tool/test_util.cc
@@ -258,11 +258,8 @@ std::string GetTestFilePath(absl::string_view relative_path) {
   return file::JoinPath(GetTestRootDir(), relative_path);
 }
 
-absl::StatusOr<std::unique_ptr<ImageFrame>> LoadTestImage(
-    absl::string_view path, ImageFormat::Format format) {
-  std::string encoded;
-  MP_RETURN_IF_ERROR(mediapipe::file::GetContents(path, &encoded));
-
+absl::StatusOr<std::unique_ptr<ImageFrame>> DecodeTestImage(
+    absl::string_view encoded, ImageFormat::Format format) {
   // stbi_load determines the output pixel format based on the desired channels.
   // 0 means "use whatever's in the file".
   int desired_channels = format == ImageFormat::UNKNOWN ? 0
@@ -274,10 +271,10 @@ absl::StatusOr<std::unique_ptr<ImageFrame>> LoadTestImage(
       << "unsupported output format requested: " << format;
 
   int width, height, channels_in_file;
-  auto data = stbi_load_from_memory(reinterpret_cast<stbi_uc*>(encoded.data()),
-                                    encoded.size(), &width, &height,
-                                    &channels_in_file, desired_channels);
-  RET_CHECK(data) << "failed to decode image data from: " << path;
+  auto data = stbi_load_from_memory(
+      reinterpret_cast<const stbi_uc*>(encoded.data()), encoded.size(), &width,
+      &height, &channels_in_file, desired_channels);
+  RET_CHECK(data) << "failed to decode image data";
 
   // If we didn't specify a desired format, it will be determined by what the
   // file contains.
@@ -295,6 +292,13 @@ absl::StatusOr<std::unique_ptr<ImageFrame>> LoadTestImage(
       format, width, height, width * output_channels, data, stbi_image_free);
 }
 
+absl::StatusOr<std::unique_ptr<ImageFrame>> LoadTestImage(
+    absl::string_view path, ImageFormat::Format format) {
+  std::string encoded;
+  MP_RETURN_IF_ERROR(mediapipe::file::GetContents(path, &encoded));
+  return DecodeTestImage(encoded, format);
+}
+
 std::unique_ptr<ImageFrame> LoadTestPng(absl::string_view path,
                                         ImageFormat::Format format) {
   return nullptr;
diff --git a/mediapipe/framework/tool/test_util.h b/mediapipe/framework/tool/test_util.h
index 71c096db7..80b768e3d 100644
--- a/mediapipe/framework/tool/test_util.h
+++ b/mediapipe/framework/tool/test_util.h
@@ -81,6 +81,10 @@ std::string GetTestDataDir(absl::string_view package_base_path);
 // Loads a binary graph from path. Returns true iff successful.
 bool LoadTestGraph(CalculatorGraphConfig* proto, const std::string& path);
 
+// Loads an image from memory.
+absl::StatusOr<std::unique_ptr<ImageFrame>> DecodeTestImage(
+    absl::string_view encoded, ImageFormat::Format format = ImageFormat::SRGBA);
+
 // Loads an image from path.
 absl::StatusOr<std::unique_ptr<ImageFrame>> LoadTestImage(
     absl::string_view path, ImageFormat::Format format = ImageFormat::SRGBA);