Add VEC32F4 support to ImageFrame
PiperOrigin-RevId: 522153305
This commit is contained in:
parent
6605f551e7
commit
065d750781
|
@ -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;
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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);
|
||||
}
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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)));
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue
Block a user