diff --git a/mediapipe/framework/formats/image_format.proto b/mediapipe/framework/formats/image_format.proto index 61e004ac6..e9b69a4c1 100644 --- a/mediapipe/framework/formats/image_format.proto +++ b/mediapipe/framework/formats/image_format.proto @@ -69,6 +69,9 @@ message ImageFormat { // Two floats per pixel. VEC32F2 = 12; + // Four floats per pixel. + VEC32F4 = 13; + // LAB, interleaved: one byte for L, then one byte for a, then one // byte for b for each pixel. LAB8 = 10; diff --git a/mediapipe/framework/formats/image_frame.cc b/mediapipe/framework/formats/image_frame.cc index 772c91014..2de819a35 100644 --- a/mediapipe/framework/formats/image_frame.cc +++ b/mediapipe/framework/formats/image_frame.cc @@ -280,6 +280,8 @@ int ImageFrame::NumberOfChannelsForFormat(ImageFormat::Format format) { return 1; case ImageFormat::VEC32F2: return 2; + case ImageFormat::VEC32F4: + return 4; case ImageFormat::LAB8: return 3; case ImageFormat::SBGRA: @@ -309,6 +311,8 @@ int ImageFrame::ChannelSizeForFormat(ImageFormat::Format format) { return sizeof(float); case ImageFormat::VEC32F2: return sizeof(float); + case ImageFormat::VEC32F4: + return sizeof(float); case ImageFormat::LAB8: return sizeof(uint8_t); case ImageFormat::SBGRA: @@ -338,6 +342,8 @@ int ImageFrame::ByteDepthForFormat(ImageFormat::Format format) { return 4; case ImageFormat::VEC32F2: return 4; + case ImageFormat::VEC32F4: + return 4; case ImageFormat::LAB8: return 1; case ImageFormat::SBGRA: diff --git a/mediapipe/framework/formats/image_frame_opencv.cc b/mediapipe/framework/formats/image_frame_opencv.cc index 940e18263..1ba8c719f 100644 --- a/mediapipe/framework/formats/image_frame_opencv.cc +++ b/mediapipe/framework/formats/image_frame_opencv.cc @@ -59,6 +59,9 @@ int GetMatType(const mediapipe::ImageFormat::Format format) { case mediapipe::ImageFormat::VEC32F2: type = CV_32FC2; break; + case mediapipe::ImageFormat::VEC32F4: + type = CV_32FC4; + break; case mediapipe::ImageFormat::LAB8: type = CV_8U; break; diff --git a/mediapipe/framework/formats/image_frame_opencv_test.cc b/mediapipe/framework/formats/image_frame_opencv_test.cc index ae6f90f81..87d2ffb36 100644 --- a/mediapipe/framework/formats/image_frame_opencv_test.cc +++ b/mediapipe/framework/formats/image_frame_opencv_test.cc @@ -113,6 +113,7 @@ TEST(ImageFrameOpencvTest, ImageFormats) { ImageFrame frame_g16(ImageFormat::GRAY16, i_width, i_height); ImageFrame frame_v32f1(ImageFormat::VEC32F1, i_width, i_height); ImageFrame frame_v32f2(ImageFormat::VEC32F2, i_width, i_height); + ImageFrame frame_v32f4(ImageFormat::VEC32F4, i_width, i_height); ImageFrame frame_c3(ImageFormat::SRGB, i_width, i_height); ImageFrame frame_c4(ImageFormat::SRGBA, i_width, i_height); @@ -120,6 +121,7 @@ TEST(ImageFrameOpencvTest, ImageFormats) { cv::Mat mat_g16 = formats::MatView(&frame_g16); cv::Mat mat_v32f1 = formats::MatView(&frame_v32f1); cv::Mat mat_v32f2 = formats::MatView(&frame_v32f2); + cv::Mat mat_v32f4 = formats::MatView(&frame_v32f4); cv::Mat mat_c3 = formats::MatView(&frame_c3); cv::Mat mat_c4 = formats::MatView(&frame_c4); @@ -127,6 +129,7 @@ TEST(ImageFrameOpencvTest, ImageFormats) { EXPECT_EQ(mat_g16.type(), CV_16UC1); EXPECT_EQ(mat_v32f1.type(), CV_32FC1); EXPECT_EQ(mat_v32f2.type(), CV_32FC2); + EXPECT_EQ(mat_v32f4.type(), CV_32FC4); EXPECT_EQ(mat_c3.type(), CV_8UC3); EXPECT_EQ(mat_c4.type(), CV_8UC4); } diff --git a/mediapipe/framework/formats/image_opencv.cc b/mediapipe/framework/formats/image_opencv.cc index 7d9ce4a13..498c7831f 100644 --- a/mediapipe/framework/formats/image_opencv.cc +++ b/mediapipe/framework/formats/image_opencv.cc @@ -60,6 +60,9 @@ int GetMatType(const mediapipe::ImageFormat::Format format) { case mediapipe::ImageFormat::VEC32F2: type = CV_32FC2; break; + case mediapipe::ImageFormat::VEC32F4: + type = CV_32FC4; + break; case mediapipe::ImageFormat::LAB8: type = CV_8U; break; diff --git a/mediapipe/framework/tool/test_util.cc b/mediapipe/framework/tool/test_util.cc index d05171d20..5642941e9 100644 --- a/mediapipe/framework/tool/test_util.cc +++ b/mediapipe/framework/tool/test_util.cc @@ -191,6 +191,7 @@ absl::Status CompareImageFrames(const ImageFrame& image1, max_alpha_diff, max_avg_diff, diff_image); case ImageFormat::VEC32F1: case ImageFormat::VEC32F2: + case ImageFormat::VEC32F4: return CompareDiff(image1, image2, max_color_diff, max_alpha_diff, max_avg_diff, diff_image); default: diff --git a/mediapipe/gpu/gpu_buffer_format.cc b/mediapipe/gpu/gpu_buffer_format.cc index 8e2e3858e..a820f04d6 100644 --- a/mediapipe/gpu/gpu_buffer_format.cc +++ b/mediapipe/gpu/gpu_buffer_format.cc @@ -204,6 +204,8 @@ ImageFormat::Format ImageFormatForGpuBufferFormat(GpuBufferFormat format) { return ImageFormat::SRGB; case GpuBufferFormat::kTwoComponentFloat32: return ImageFormat::VEC32F2; + case GpuBufferFormat::kRGBAFloat128: + return ImageFormat::VEC32F4; case GpuBufferFormat::kRGBA32: // TODO: this likely maps to ImageFormat::SRGBA case GpuBufferFormat::kGrayHalf16: @@ -211,7 +213,6 @@ ImageFormat::Format ImageFormatForGpuBufferFormat(GpuBufferFormat format) { case GpuBufferFormat::kTwoComponent8: case GpuBufferFormat::kTwoComponentHalf16: case GpuBufferFormat::kRGBAHalf64: - case GpuBufferFormat::kRGBAFloat128: case GpuBufferFormat::kNV12: case GpuBufferFormat::kNV21: case GpuBufferFormat::kI420: @@ -232,6 +233,8 @@ GpuBufferFormat GpuBufferFormatForImageFormat(ImageFormat::Format format) { return GpuBufferFormat::kGrayFloat32; case ImageFormat::VEC32F2: return GpuBufferFormat::kTwoComponentFloat32; + case ImageFormat::VEC32F4: + return GpuBufferFormat::kRGBAFloat128; case ImageFormat::GRAY8: return GpuBufferFormat::kOneComponent8; case ImageFormat::YCBCR420P: diff --git a/mediapipe/objc/util.cc b/mediapipe/objc/util.cc index 895463060..36ad4e195 100644 --- a/mediapipe/objc/util.cc +++ b/mediapipe/objc/util.cc @@ -365,6 +365,10 @@ absl::StatusOr> CreateCVPixelBufferForImageFrame( pixel_format = kCVPixelFormatType_TwoComponent32Float; break; + case mediapipe::ImageFormat::VEC32F4: + pixel_format = kCVPixelFormatType_128RGBAFloat; + break; + default: return ::mediapipe::UnknownErrorBuilder(MEDIAPIPE_LOC) << "unsupported ImageFrame format: " << image_format; @@ -440,6 +444,10 @@ absl::StatusOr> CreateCVPixelBufferCopyingImageFrame( pixel_format = kCVPixelFormatType_TwoComponent32Float; break; + case mediapipe::ImageFormat::VEC32F4: + pixel_format = kCVPixelFormatType_128RGBAFloat; + break; + default: return ::mediapipe::UnknownErrorBuilder(MEDIAPIPE_LOC) << "unsupported ImageFrame format: " << image_format; diff --git a/mediapipe/python/pybind/image.cc b/mediapipe/python/pybind/image.cc index 6049abfae..5e5ba7530 100644 --- a/mediapipe/python/pybind/image.cc +++ b/mediapipe/python/pybind/image.cc @@ -120,16 +120,17 @@ void ImageSubmodule(pybind11::module* module) { py::init([](mediapipe::ImageFormat::Format format, const py::array_t& data) { if (format != mediapipe::ImageFormat::VEC32F1 && - format != mediapipe::ImageFormat::VEC32F2) { + format != mediapipe::ImageFormat::VEC32F2 && + format != mediapipe::ImageFormat::VEC32F4) { throw RaisePyError( PyExc_RuntimeError, - "float image data should be either VEC32F1 or VEC32F2 " - "MediaPipe image formats."); + "float image data should be either VEC32F1, VEC32F2, or " + "VEC32F4 MediaPipe image formats."); } return Image(std::shared_ptr( CreateImageFrame(format, data))); }), - R"doc(For float data type, valid ImageFormat are VEC32F1 and VEC32F2.)doc", + R"doc(For float data type, valid ImageFormat are VEC32F1, VEC32F2, and VEC32F4.)doc", py::arg("image_format"), py::arg("data").noconvert()); image.def( diff --git a/mediapipe/python/pybind/packet_creator.cc b/mediapipe/python/pybind/packet_creator.cc index 92e695020..c8ae7c259 100644 --- a/mediapipe/python/pybind/packet_creator.cc +++ b/mediapipe/python/pybind/packet_creator.cc @@ -42,7 +42,8 @@ Packet CreateImageFramePacket(mediapipe::ImageFormat::Format format, format == mediapipe::ImageFormat::SRGBA64) { return Adopt(CreateImageFrame(format, data, copy).release()); } else if (format == mediapipe::ImageFormat::VEC32F1 || - format == mediapipe::ImageFormat::VEC32F2) { + format == mediapipe::ImageFormat::VEC32F2 || + format == mediapipe::ImageFormat::VEC32F4) { return Adopt(CreateImageFrame(format, data, copy).release()); } throw RaisePyError(PyExc_RuntimeError, @@ -63,7 +64,8 @@ Packet CreateImagePacket(mediapipe::ImageFormat::Format format, return MakePacket(std::shared_ptr( CreateImageFrame(format, data, copy))); } else if (format == mediapipe::ImageFormat::VEC32F1 || - format == mediapipe::ImageFormat::VEC32F2) { + format == mediapipe::ImageFormat::VEC32F2 || + format == mediapipe::ImageFormat::VEC32F4) { return MakePacket(std::shared_ptr( CreateImageFrame(format, data, copy))); }