From 5998e96eed8e6ee43de445013be82a7f482ac0f0 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Thu, 23 Mar 2023 03:31:45 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 518814155 --- mediapipe/util/frame_buffer/BUILD | 5 + mediapipe/util/frame_buffer/float_buffer.cc | 57 +++++++++++ mediapipe/util/frame_buffer/float_buffer.h | 85 ++++++++++++++++ .../util/frame_buffer/frame_buffer_util.cc | 62 ++++++++++++ .../util/frame_buffer/frame_buffer_util.h | 9 ++ .../frame_buffer/frame_buffer_util_test.cc | 22 +++++ mediapipe/util/frame_buffer/halide/BUILD | 6 ++ .../halide/rgb_float_generator.cc | 59 ++++++++++++ mediapipe/util/frame_buffer/rgb_buffer.cc | 8 ++ mediapipe/util/frame_buffer/rgb_buffer.h | 4 + .../util/frame_buffer/rgb_buffer_test.cc | 96 +++++++++++++++++++ 11 files changed, 413 insertions(+) create mode 100644 mediapipe/util/frame_buffer/float_buffer.cc create mode 100644 mediapipe/util/frame_buffer/float_buffer.h create mode 100644 mediapipe/util/frame_buffer/halide/rgb_float_generator.cc diff --git a/mediapipe/util/frame_buffer/BUILD b/mediapipe/util/frame_buffer/BUILD index 27343d6df..5dfffbac7 100644 --- a/mediapipe/util/frame_buffer/BUILD +++ b/mediapipe/util/frame_buffer/BUILD @@ -22,6 +22,7 @@ cc_library( deps = [ ":buffer", "//mediapipe/framework/formats:frame_buffer", + "//mediapipe/framework/formats:tensor", "//mediapipe/framework/port:status", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", @@ -37,6 +38,7 @@ cc_test( deps = [ ":frame_buffer_util", "//mediapipe/framework/formats:frame_buffer", + "//mediapipe/framework/formats:tensor", "//mediapipe/framework/port:gtest_main", "//mediapipe/framework/port:status", ], @@ -46,12 +48,14 @@ cc_library( name = "buffer", srcs = [ "buffer_common.cc", + "float_buffer.cc", "gray_buffer.cc", "rgb_buffer.cc", "yuv_buffer.cc", ], hdrs = [ "buffer_common.h", + "float_buffer.h", "gray_buffer.h", "rgb_buffer.h", "yuv_buffer.h", @@ -61,6 +65,7 @@ cc_library( "//mediapipe/util/frame_buffer/halide:gray_resize_halide", "//mediapipe/util/frame_buffer/halide:gray_rotate_halide", "//mediapipe/util/frame_buffer/halide:rgb_flip_halide", + "//mediapipe/util/frame_buffer/halide:rgb_float_halide", "//mediapipe/util/frame_buffer/halide:rgb_gray_halide", "//mediapipe/util/frame_buffer/halide:rgb_resize_halide", "//mediapipe/util/frame_buffer/halide:rgb_rgb_halide", diff --git a/mediapipe/util/frame_buffer/float_buffer.cc b/mediapipe/util/frame_buffer/float_buffer.cc new file mode 100644 index 000000000..f5772abe3 --- /dev/null +++ b/mediapipe/util/frame_buffer/float_buffer.cc @@ -0,0 +1,57 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mediapipe/util/frame_buffer/float_buffer.h" + +#include + +namespace mediapipe { +namespace frame_buffer { + +FloatBuffer::FloatBuffer(float* data, int width, int height, int channels) + : owned_buffer_(nullptr) { + Initialize(data, width, height, channels); +} + +FloatBuffer::FloatBuffer(int width, int height, int channels) { + owned_buffer_ = std::make_unique(FloatSize(width, height, channels)); + Initialize(owned_buffer_.get(), width, height, channels); +} + +FloatBuffer::FloatBuffer(const FloatBuffer& other) : buffer_(other.buffer_) {} + +FloatBuffer::FloatBuffer(FloatBuffer&& other) { *this = std::move(other); } + +FloatBuffer& FloatBuffer::operator=(const FloatBuffer& other) { + if (this != &other) { + buffer_ = other.buffer_; + } + return *this; +} +FloatBuffer& FloatBuffer::operator=(FloatBuffer&& other) { + if (this != &other) { + buffer_ = other.buffer_; + } + return *this; +} + +FloatBuffer::~FloatBuffer() {} + +void FloatBuffer::Initialize(float* data, int width, int height, int channels) { + buffer_ = Halide::Runtime::Buffer::make_interleaved(data, width, + height, channels); +} + +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/float_buffer.h b/mediapipe/util/frame_buffer/float_buffer.h new file mode 100644 index 000000000..9dc723789 --- /dev/null +++ b/mediapipe/util/frame_buffer/float_buffer.h @@ -0,0 +1,85 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MEDIAPIPE_UTIL_FRAME_BUFFER_FLOAT_BUFFER_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_FLOAT_BUFFER_H_ + +#include "HalideBuffer.h" +#include "HalideRuntime.h" + +namespace mediapipe { +namespace frame_buffer { + +// FloatBuffer represents a view over an interleaved floating-point image. +// +// FloatBuffers may be copied and moved efficiently; their backing buffers are +// shared and never deep copied. +// +// FloatBuffer requires a minimum image width depending on the natural vector +// size of the platform, e.g., 16px. This is not validated by FloatBuffer. +class FloatBuffer { + public: + // Returns the size (in number of float) of a FloatBuffer given dimensions. + static int FloatSize(int width, int height, int channels) { + return width * height * channels; + } + + // Builds a FloatBuffer using the given backing buffer and dimensions. + FloatBuffer(float* data, int width, int height, int channels); + + // Builds a FloatBuffer using the given dimensions. + // + // The underlying backing buffer if allocated and owned by this FloatBuffer. + FloatBuffer(int width, int height, int channels); + + // FloatBuffer is copyable. The source retains ownership of its backing + // buffer. + FloatBuffer(const FloatBuffer& other); + // FloatBuffer is moveable. The source loses ownership of any backing buffers. + FloatBuffer(FloatBuffer&& other); + // FloatBuffer is assignable. + FloatBuffer& operator=(const FloatBuffer& other); + FloatBuffer& operator=(FloatBuffer&& other); + + ~FloatBuffer(); + + // Release ownership of the owned backing buffer. + float* Release() { return owned_buffer_.release(); } + + // Returns the halide_buffer_t* for the image. + const halide_buffer_t* buffer() const { return buffer_.raw_buffer(); } + // Returns the halide_buffer_t* for the image. + halide_buffer_t* buffer() { return buffer_.raw_buffer(); } + + // Returns the image width. + int width() const { return buffer_.dim(0).extent(); } + // Returns the image height. + int height() const { return buffer_.dim(1).extent(); } + // Returns the number of channels. + int channels() const { return buffer_.dim(2).extent(); } + + private: + void Initialize(float* data, int width, int height, int channels); + + // Non-NULL iff this FloatBuffer owns its backing buffer. + std::unique_ptr owned_buffer_; + + // Backing buffer: layout is always width x height x channel (interleaved). + Halide::Runtime::Buffer buffer_; +}; + +} // namespace frame_buffer +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_FLOAT_BUFFER_H_ diff --git a/mediapipe/util/frame_buffer/frame_buffer_util.cc b/mediapipe/util/frame_buffer/frame_buffer_util.cc index b18e8cb13..fa0ea9094 100644 --- a/mediapipe/util/frame_buffer/frame_buffer_util.cc +++ b/mediapipe/util/frame_buffer/frame_buffer_util.cc @@ -22,7 +22,9 @@ #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/tensor.h" #include "mediapipe/framework/port/status_macros.h" +#include "mediapipe/util/frame_buffer/float_buffer.h" #include "mediapipe/util/frame_buffer/gray_buffer.h" #include "mediapipe/util/frame_buffer/rgb_buffer.h" #include "mediapipe/util/frame_buffer/yuv_buffer.h" @@ -50,6 +52,22 @@ bool IsSupportedYuvBuffer(const FrameBuffer& buffer) { buffer.format() == FrameBuffer::Format::kYV21; } +// Returns the number of channels for the provided buffer. Returns an error if +// the buffer is not using an interleaved single-planar format. +absl::StatusOr NumberOfChannels(const FrameBuffer& buffer) { + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + return kGrayChannel; + case FrameBuffer::Format::kRGB: + return kRgbChannels; + case FrameBuffer::Format::kRGBA: + return kRgbaChannels; + default: + return absl::InvalidArgumentError( + absl::StrFormat("Unsupported buffer format: %i.", buffer.format())); + } +} + // Shared validation functions. //------------------------------------------------------------------------------ @@ -216,6 +234,25 @@ absl::Status ValidateConvertFormats(FrameBuffer::Format from_format, } } +absl::Status ValidateFloatTensorInputs(const FrameBuffer& buffer, + const Tensor& tensor) { + if (tensor.element_type() != Tensor::ElementType::kFloat32) { + return absl::InvalidArgumentError(absl::StrFormat( + "Tensor type %i is not supported.", tensor.element_type())); + } + const auto& shape = tensor.shape(); + if (shape.dims.size() != 4 || shape.dims[0] != 1) { + return absl::InvalidArgumentError("Expected tensor with batch size of 1."); + } + ASSIGN_OR_RETURN(int channels, NumberOfChannels(buffer)); + if (shape.dims[2] != buffer.dimension().width || + shape.dims[1] != buffer.dimension().height || shape.dims[3] != channels) { + return absl::InvalidArgumentError( + "Input buffer and output tensor must have the same dimensions."); + } + return absl::OkStatus(); +} + // Construct buffer helper functions. //------------------------------------------------------------------------------ @@ -380,6 +417,19 @@ absl::Status RotateRgb(const FrameBuffer& buffer, int angle, : absl::UnknownError("Halide rgb[a] rotate operation failed."); } +absl::Status ToFloatTensorRgb(const FrameBuffer& buffer, float scale, + float offset, Tensor& tensor) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + ASSIGN_OR_RETURN(int channels, NumberOfChannels(buffer)); + auto view = tensor.GetCpuWriteView(); + float* data = view.buffer(); + FloatBuffer output(data, buffer.dimension().width, buffer.dimension().height, + channels); + return input.ToFloat(scale, offset, &output) + ? absl::OkStatus() + : absl::UnknownError("Halide rgb[a] to float conversion failed."); +} + // Yuv transformation functions. //------------------------------------------------------------------------------ @@ -717,6 +767,18 @@ absl::Status Convert(const FrameBuffer& buffer, FrameBuffer* output_buffer) { } } +absl::Status ToFloatTensor(const FrameBuffer& buffer, float scale, float offset, + Tensor& tensor) { + MP_RETURN_IF_ERROR(ValidateFloatTensorInputs(buffer, tensor)); + switch (buffer.format()) { + case FrameBuffer::Format::kRGB: + return ToFloatTensorRgb(buffer, scale, offset, tensor); + default: + return absl::InvalidArgumentError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + int GetFrameBufferByteSize(FrameBuffer::Dimension dimension, FrameBuffer::Format format) { switch (format) { diff --git a/mediapipe/util/frame_buffer/frame_buffer_util.h b/mediapipe/util/frame_buffer/frame_buffer_util.h index eb17e0094..ca289f185 100644 --- a/mediapipe/util/frame_buffer/frame_buffer_util.h +++ b/mediapipe/util/frame_buffer/frame_buffer_util.h @@ -21,6 +21,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" #include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/tensor.h" namespace mediapipe { namespace frame_buffer { @@ -98,6 +99,14 @@ absl::Status FlipVertically(const FrameBuffer& buffer, // on the buffer and output_buffer dimensions. absl::Status Convert(const FrameBuffer& buffer, FrameBuffer* output_buffer); +// Converts `buffer` into the provided float Tensor. Each value is converted to +// a float using: +// output = input * scale + offset +// +// Note that only interleaved single-planar formats support this operation. +absl::Status ToFloatTensor(const FrameBuffer& buffer, float scale, float offset, + Tensor& tensor); + // Miscellaneous Methods // ----------------------------------------------------------------- diff --git a/mediapipe/util/frame_buffer/frame_buffer_util_test.cc b/mediapipe/util/frame_buffer/frame_buffer_util_test.cc index d92eb3f53..8e86f02d0 100644 --- a/mediapipe/util/frame_buffer/frame_buffer_util_test.cc +++ b/mediapipe/util/frame_buffer/frame_buffer_util_test.cc @@ -19,6 +19,7 @@ #include #include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/tensor.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" #include "mediapipe/framework/port/status_macros.h" @@ -520,6 +521,27 @@ TEST(FrameBufferUtil, RgbaToRgbConversion) { EXPECT_EQ(output_data[5], 100); } +TEST(FrameBufferUtil, RgbToFloatTensor) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 2, .height = 1}; + constexpr float kScale = 0.1f, kOffset = 0.1f; + uint8_t data[] = {1, 2, 3, 4, 5, 6}; + auto input = CreateFromRgbRawBuffer(data, kBufferDimension); + Tensor output( + Tensor::ElementType::kFloat32, + Tensor::Shape{1, kBufferDimension.height, kBufferDimension.width, 3}); + + MP_ASSERT_OK(ToFloatTensor(*input, kScale, kOffset, output)); + + auto view = output.GetCpuReadView(); + const float* output_data = view.buffer(); + EXPECT_EQ(output_data[0], 0.2f); + EXPECT_EQ(output_data[1], 0.3f); + EXPECT_EQ(output_data[2], 0.4f); + EXPECT_EQ(output_data[3], 0.5f); + EXPECT_EQ(output_data[4], 0.6f); + EXPECT_EQ(output_data[5], 0.7f); +} + TEST(FrameBufferUtil, RgbaCrop) { constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, kOutputDimension = {.width = 1, .height = 1}; diff --git a/mediapipe/util/frame_buffer/halide/BUILD b/mediapipe/util/frame_buffer/halide/BUILD index 619a16d26..b6cf3c7ed 100644 --- a/mediapipe/util/frame_buffer/halide/BUILD +++ b/mediapipe/util/frame_buffer/halide/BUILD @@ -62,6 +62,12 @@ halide_library( generator_name = "rgb_rgb_generator", ) +halide_library( + name = "rgb_float_halide", + srcs = ["rgb_float_generator.cc"], + generator_name = "rgb_float_generator", +) + # YUV operations: halide_library( name = "yuv_flip_halide", diff --git a/mediapipe/util/frame_buffer/halide/rgb_float_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_float_generator.cc new file mode 100644 index 000000000..994f731b0 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_float_generator.cc @@ -0,0 +1,59 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "Halide.h" + +namespace { + +class RgbFloat : public Halide::Generator { + public: + Var x{"x"}, y{"y"}, c{"c"}; + + Input> src_rgb{"src_rgb"}; + Input scale{"scale"}; + Input offset{"offset"}; + Output> dst_float{"dst_float"}; + + void generate(); + void schedule(); +}; + +void RgbFloat::generate() { + dst_float(x, y, c) = Halide::cast(src_rgb(x, y, c)) * scale + offset; +} + +void RgbFloat::schedule() { + Halide::Expr input_rgb_channels = src_rgb.dim(2).extent(); + Halide::Expr output_float_channels = dst_float.dim(2).extent(); + + // The source buffer starts at zero in every dimension and requires an + // interleaved format. + src_rgb.dim(0).set_min(0); + src_rgb.dim(1).set_min(0); + src_rgb.dim(2).set_min(0); + src_rgb.dim(0).set_stride(input_rgb_channels); + src_rgb.dim(2).set_stride(1); + + // The destination buffer starts at zero in every dimension and requires an + // interleaved format. + dst_float.dim(0).set_min(0); + dst_float.dim(1).set_min(0); + dst_float.dim(2).set_min(0); + dst_float.dim(0).set_stride(output_float_channels); + dst_float.dim(2).set_stride(1); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbFloat, rgb_float_generator) diff --git a/mediapipe/util/frame_buffer/rgb_buffer.cc b/mediapipe/util/frame_buffer/rgb_buffer.cc index 9ae849eab..432e32e9e 100644 --- a/mediapipe/util/frame_buffer/rgb_buffer.cc +++ b/mediapipe/util/frame_buffer/rgb_buffer.cc @@ -17,8 +17,10 @@ #include #include "mediapipe/util/frame_buffer/buffer_common.h" +#include "mediapipe/util/frame_buffer/float_buffer.h" #include "mediapipe/util/frame_buffer/gray_buffer.h" #include "mediapipe/util/frame_buffer/halide/rgb_flip_halide.h" +#include "mediapipe/util/frame_buffer/halide/rgb_float_halide.h" #include "mediapipe/util/frame_buffer/halide/rgb_gray_halide.h" #include "mediapipe/util/frame_buffer/halide/rgb_resize_halide.h" #include "mediapipe/util/frame_buffer/halide/rgb_rgb_halide.h" @@ -122,6 +124,12 @@ bool RgbBuffer::Convert(RgbBuffer* output) { return result == 0; } +bool RgbBuffer::ToFloat(float scale, float offset, FloatBuffer* output) { + const int result = + rgb_float_halide(buffer(), scale, offset, output->buffer()); + return result == 0; +} + void RgbBuffer::Initialize(uint8_t* data, int width, int height, bool alpha) { const int channels = alpha ? 4 : 3; buffer_ = Halide::Runtime::Buffer::make_interleaved( diff --git a/mediapipe/util/frame_buffer/rgb_buffer.h b/mediapipe/util/frame_buffer/rgb_buffer.h index 06423c3f6..57478775c 100644 --- a/mediapipe/util/frame_buffer/rgb_buffer.h +++ b/mediapipe/util/frame_buffer/rgb_buffer.h @@ -19,6 +19,7 @@ #include "HalideBuffer.h" #include "HalideRuntime.h" +#include "mediapipe/util/frame_buffer/float_buffer.h" #include "mediapipe/util/frame_buffer/gray_buffer.h" #include "mediapipe/util/frame_buffer/yuv_buffer.h" @@ -106,6 +107,9 @@ class RgbBuffer { // Performs a rgb to rgba / rgba to rgb format conversion. bool Convert(RgbBuffer* output); + // Performs a RGB to float conversion. + bool ToFloat(float scale, float offset, FloatBuffer* output); + // Release ownership of the owned backing buffer. uint8_t* Release() { return owned_buffer_.release(); } diff --git a/mediapipe/util/frame_buffer/rgb_buffer_test.cc b/mediapipe/util/frame_buffer/rgb_buffer_test.cc index e5cb39c69..8ade0b927 100644 --- a/mediapipe/util/frame_buffer/rgb_buffer_test.cc +++ b/mediapipe/util/frame_buffer/rgb_buffer_test.cc @@ -14,11 +14,13 @@ #include "mediapipe/util/frame_buffer/rgb_buffer.h" +#include #include #include "absl/log/log.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" +#include "mediapipe/util/frame_buffer/float_buffer.h" #include "mediapipe/util/frame_buffer/gray_buffer.h" #include "mediapipe/util/frame_buffer/yuv_buffer.h" @@ -87,6 +89,22 @@ bool CompareArray(const uint8_t* lhs_ptr, const uint8_t* rhs_ptr, int width, return true; } +// Returns true if the data in the two arrays are the same. Otherwise, return +// false. +bool CompareArray(const float* lhs_ptr, const float* rhs_ptr, int width, + int height) { + constexpr float kTolerancy = 1e-6; + for (int i = 0; i < height; ++i) { + for (int j = 0; j < width; ++j) { + if (std::abs(lhs_ptr[i * width + j] - rhs_ptr[i * width + j]) > + kTolerancy) { + return false; + } + } + } + return true; +} + // Returns true if the halide buffers of two input GrayBuffer are identical. // Otherwise, returns false; bool CompareBuffer(const GrayBuffer& lhs, const GrayBuffer& rhs) { @@ -129,6 +147,20 @@ bool CompareBuffer(const YuvBuffer& lhs, const YuvBuffer& rhs) { lhs.height() / 2); } +// Returns true if the halide buffers of two input FloatBuffer are identical. +// Otherwise, returns false; +bool CompareBuffer(const FloatBuffer& lhs, const FloatBuffer& rhs) { + if (lhs.width() != rhs.width() || lhs.height() != rhs.height() || + lhs.channels() != rhs.channels()) { + return false; + } + const float* reference_ptr = reinterpret_cast( + const_cast(lhs).buffer()->host); + const float* converted_ptr = reinterpret_cast( + const_cast(rhs).buffer()->host); + return CompareArray(reference_ptr, converted_ptr, lhs.width(), lhs.height()); +} + TEST(RgbBufferTest, Properties) { RgbBuffer rgb(2, 8, false), rgba(2, 8, true); EXPECT_EQ(2, rgb.width()); @@ -601,6 +633,70 @@ TEST(RgbBufferTest, PaddedRgbaConvertRgb) { RgbBuffer rgb_buffer = RgbBuffer(rgb_data, kWidth, kHeight, false); EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); } + +TEST(RgbBufferTest, RgbToFloat) { + constexpr int kWidth = 2, kHeight = 1, kChannels = 3; + constexpr float kScale = 0.01f, kOffset = 0.5f; + uint8_t rgb_data[] = {200, 100, 50, 100, 50, 20}; + RgbBuffer source(rgb_data, kWidth, kHeight, false); + FloatBuffer result(kWidth, kHeight, kChannels); + + ASSERT_TRUE(source.ToFloat(kScale, kOffset, &result)); + + float float_data[] = {2.5f, 1.5f, 1.0f, 1.5f, 1.0f, 0.7f}; + FloatBuffer float_buffer = + FloatBuffer(float_data, kWidth, kHeight, kChannels); + EXPECT_TRUE(CompareBuffer(float_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbToFloat) { + constexpr int kWidth = 4, kHeight = 2, kChannels = 3; + constexpr float kScale = 0.01f, kOffset = 0.0f; + RgbBuffer source = GetPaddedRgbBuffer(); + FloatBuffer result(kWidth, kHeight, kChannels); + + ASSERT_TRUE(source.ToFloat(kScale, kOffset, &result)); + + float float_data[] = {0.1f, 0.2f, 0.3f, 0.2f, 0.3f, 0.4f, 0.3f, 0.4f, + 0.5f, 0.4f, 0.5f, 0.6f, 0.2f, 0.4f, 0.6f, 0.4f, + 0.6f, 0.8f, 0.6f, 0.8f, 1.0f, 0.8f, 1.0f, 1.2f}; + FloatBuffer float_buffer = + FloatBuffer(float_data, kWidth, kHeight, kChannels); + EXPECT_TRUE(CompareBuffer(float_buffer, result)); +} + +TEST(RgbBufferTest, RgbaToFloat) { + constexpr int kWidth = 2, kHeight = 1, kChannels = 4; + constexpr float kScale = 0.01f, kOffset = 0.5f; + uint8_t rgba_data[] = {200, 100, 50, 30, 100, 50, 20, 70}; + RgbBuffer source(rgba_data, kWidth, kHeight, true); + FloatBuffer result(kWidth, kHeight, kChannels); + + ASSERT_TRUE(source.ToFloat(kScale, kOffset, &result)); + + float float_data[] = {2.5f, 1.5f, 1.0f, 0.8f, 1.5f, 1.0f, 0.7f, 1.2f}; + FloatBuffer float_buffer = + FloatBuffer(float_data, kWidth, kHeight, kChannels); + EXPECT_TRUE(CompareBuffer(float_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaToFloat) { + constexpr int kWidth = 4, kHeight = 2, kChannels = 4; + constexpr float kScale = 0.01f, kOffset = 0.0f; + RgbBuffer source = GetPaddedRgbaBuffer(); + FloatBuffer result(kWidth, kHeight, kChannels); + + ASSERT_TRUE(source.ToFloat(kScale, kOffset, &result)); + + float float_data[] = {0.1f, 0.2f, 0.3f, 2.55f, 0.2f, 0.3f, 0.4f, 2.55f, + 0.3f, 0.4f, 0.5f, 2.55f, 0.4f, 0.5f, 0.6f, 2.55f, + 0.2f, 0.4f, 0.6f, 2.55f, 0.4f, 0.6f, 0.8f, 2.55f, + 0.6f, 0.8f, 1.0f, 2.55f, 0.8f, 1.0f, 1.2f, 2.55f}; + FloatBuffer float_buffer = + FloatBuffer(float_data, kWidth, kHeight, kChannels); + EXPECT_TRUE(CompareBuffer(float_buffer, result)); +} + } // namespace } // namespace frame_buffer } // namespace mediapipe