Add VEC32F4 support to ImageFrame

PiperOrigin-RevId: 522153305

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;

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:

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;

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);
 }

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;

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<float>(image1, image2, max_color_diff, max_alpha_diff,
                                 max_avg_diff, diff_image);
     default:

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:

mediapipe/objc/util.cc

@@ -365,6 +365,10 @@ absl::StatusOr<CFHolder<CVPixelBufferRef>> 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<CFHolder<CVPixelBufferRef>> 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;

mediapipe/python/pybind/image.cc

@@ -120,16 +120,17 @@ void ImageSubmodule(pybind11::module* module) {
           py::init([](mediapipe::ImageFormat::Format format,
                       const py::array_t<float, py::array::c_style>& 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 "
+                  "float image data should be either VEC32F1, VEC32F2, or "
-                  "MediaPipe image formats.");
+                  "VEC32F4 MediaPipe image formats.");
             }
             return Image(std::shared_ptr<ImageFrame>(
                 CreateImageFrame<float>(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(

mediapipe/python/pybind/packet_creator.cc

@@ -42,7 +42,8 @@ Packet CreateImageFramePacket(mediapipe::ImageFormat::Format format,
              format == mediapipe::ImageFormat::SRGBA64) {
     return Adopt(CreateImageFrame<uint16_t>(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<float>(format, data, copy).release());
   }
   throw RaisePyError(PyExc_RuntimeError,
@@ -63,7 +64,8 @@ Packet CreateImagePacket(mediapipe::ImageFormat::Format format,
     return MakePacket<Image>(std::shared_ptr<ImageFrame>(
         CreateImageFrame<uint16_t>(format, data, copy)));
   } else if (format == mediapipe::ImageFormat::VEC32F1 ||
-             format == mediapipe::ImageFormat::VEC32F2) {
+             format == mediapipe::ImageFormat::VEC32F2 ||
+             format == mediapipe::ImageFormat::VEC32F4) {
     return MakePacket<Image>(std::shared_ptr<ImageFrame>(
         CreateImageFrame<float>(format, data, copy)));
   }