From 7fdbbee5be85c06758c01211ea0c81bbcf223ccc Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Tue, 2 May 2023 09:04:56 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 528799585
---
 mediapipe/framework/tool/BUILD        | 1 +
 mediapipe/framework/tool/test_util.cc | 8 ++++++++
 2 files changed, 9 insertions(+)

diff --git a/mediapipe/framework/tool/BUILD b/mediapipe/framework/tool/BUILD
index fbdcf8c9e..4ae0bb607 100644
--- a/mediapipe/framework/tool/BUILD
+++ b/mediapipe/framework/tool/BUILD
@@ -791,6 +791,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@stblib//:stb_image",
         "@stblib//:stb_image_write",
     ],
diff --git a/mediapipe/framework/tool/test_util.cc b/mediapipe/framework/tool/test_util.cc
index 5642941e9..64b5072c5 100644
--- a/mediapipe/framework/tool/test_util.cc
+++ b/mediapipe/framework/tool/test_util.cc
@@ -26,6 +26,7 @@
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/substitute.h"
 #include "mediapipe/framework/calculator.pb.h"
@@ -311,6 +312,13 @@ std::unique_ptr<ImageFrame> LoadTestPng(absl::string_view path,
 // Returns the path to the output if successful.
 absl::StatusOr<std::string> SavePngTestOutput(
     const mediapipe::ImageFrame& image, absl::string_view prefix) {
+  absl::flat_hash_set<ImageFormat::Format> supported_formats = {
+      ImageFormat::GRAY8, ImageFormat::SRGB, ImageFormat::SRGBA,
+      ImageFormat::LAB8, ImageFormat::SBGRA};
+  if (!supported_formats.contains(image.Format())) {
+    return absl::CancelledError(
+        absl::StrFormat("Format %d can not be saved to PNG.", image.Format()));
+  }
   std::string now_string = absl::FormatTime(absl::Now());
   std::string output_relative_path =
       absl::StrCat(prefix, "_", now_string, ".png");