cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add VEC32F4 support to ImageFrame

Browse Source 
PiperOrigin-RevId: 522153305
This commit is contained in:
MediaPipe Team 2023-04-05 14:22:26 -07:00 committed by Copybara-Service
parent 6605f551e7
commit 065d750781
10 changed files with 40 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mediapipe/framework/formats/image_format.proto
View File

@ -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;

6
mediapipe/framework/formats/image_frame.cc
View File

@ -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:

3
mediapipe/framework/formats/image_frame_opencv.cc
View File

@ -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;

3
mediapipe/framework/formats/image_frame_opencv_test.cc
View File

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

3
mediapipe/framework/formats/image_opencv.cc
View File

@ -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;

1
mediapipe/framework/tool/test_util.cc
View File

@ -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:

5
mediapipe/gpu/gpu_buffer_format.cc
View File

@ -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:

8
mediapipe/objc/util.cc
View File

@ -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;

9
mediapipe/python/pybind/image.cc
View File

@ -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 "
"MediaPipe image formats.");
"float image data should be either VEC32F1, VEC32F2, or "
"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(

6
mediapipe/python/pybind/packet_creator.cc
View File

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