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 LoadTestPng(absl::string_view path, // Returns the path to the output if successful. absl::StatusOr SavePngTestOutput( const mediapipe::ImageFrame& image, absl::string_view prefix) { + absl::flat_hash_set 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");