From 39e2c8351fc6cb6609ca5e023980c36a9d98a6c9 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Thu, 9 Mar 2023 05:04:18 -0800 Subject: [PATCH] Add build system for Halide and expose FrameBufferUtils. PiperOrigin-RevId: 515304264 --- WORKSPACE | 32 + mediapipe/util/frame_buffer/BUILD | 106 ++ mediapipe/util/frame_buffer/buffer_common.cc | 40 + mediapipe/util/frame_buffer/buffer_common.h | 32 + .../util/frame_buffer/frame_buffer_util.cc | 794 +++++++++++ .../util/frame_buffer/frame_buffer_util.h | 131 ++ .../frame_buffer/frame_buffer_util_test.cc | 1176 +++++++++++++++++ mediapipe/util/frame_buffer/gray_buffer.cc | 95 ++ mediapipe/util/frame_buffer/gray_buffer.h | 137 ++ .../util/frame_buffer/gray_buffer_test.cc | 223 ++++ mediapipe/util/frame_buffer/halide/BUILD | 118 ++ mediapipe/util/frame_buffer/halide/common.cc | 89 ++ mediapipe/util/frame_buffer/halide/common.h | 61 + .../halide/gray_flip_generator.cc | 61 + .../halide/gray_resize_generator.cc | 60 + .../halide/gray_rotate_generator.cc | 63 + .../frame_buffer/halide/rgb_flip_generator.cc | 84 ++ .../frame_buffer/halide/rgb_gray_generator.cc | 65 + .../halide/rgb_resize_generator.cc | 85 ++ .../frame_buffer/halide/rgb_rgb_generator.cc | 64 + .../halide/rgb_rotate_generator.cc | 76 ++ .../frame_buffer/halide/rgb_yuv_generator.cc | 101 ++ .../frame_buffer/halide/yuv_flip_generator.cc | 90 ++ .../halide/yuv_resize_generator.cc | 91 ++ .../frame_buffer/halide/yuv_rgb_generator.cc | 110 ++ .../halide/yuv_rotate_generator.cc | 91 ++ mediapipe/util/frame_buffer/rgb_buffer.cc | 132 ++ mediapipe/util/frame_buffer/rgb_buffer.h | 139 ++ .../util/frame_buffer/rgb_buffer_test.cc | 606 +++++++++ mediapipe/util/frame_buffer/yuv_buffer.cc | 152 +++ mediapipe/util/frame_buffer/yuv_buffer.h | 160 +++ .../util/frame_buffer/yuv_buffer_test.cc | 251 ++++ third_party/halide.BUILD | 70 + third_party/halide/BUILD | 1 + third_party/halide/BUILD.bazel | 40 + third_party/halide/halide.bzl | 875 ++++++++++++ 36 files changed, 6501 insertions(+) create mode 100644 mediapipe/util/frame_buffer/BUILD create mode 100644 mediapipe/util/frame_buffer/buffer_common.cc create mode 100644 mediapipe/util/frame_buffer/buffer_common.h create mode 100644 mediapipe/util/frame_buffer/frame_buffer_util.cc create mode 100644 mediapipe/util/frame_buffer/frame_buffer_util.h create mode 100644 mediapipe/util/frame_buffer/frame_buffer_util_test.cc create mode 100644 mediapipe/util/frame_buffer/gray_buffer.cc create mode 100644 mediapipe/util/frame_buffer/gray_buffer.h create mode 100644 mediapipe/util/frame_buffer/gray_buffer_test.cc create mode 100644 mediapipe/util/frame_buffer/halide/BUILD create mode 100644 mediapipe/util/frame_buffer/halide/common.cc create mode 100644 mediapipe/util/frame_buffer/halide/common.h create mode 100644 mediapipe/util/frame_buffer/halide/gray_flip_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/gray_resize_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/gray_rotate_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/rgb_flip_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/rgb_gray_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/rgb_resize_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/rgb_rgb_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/rgb_rotate_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/rgb_yuv_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/yuv_flip_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/yuv_resize_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/yuv_rgb_generator.cc create mode 100644 mediapipe/util/frame_buffer/halide/yuv_rotate_generator.cc create mode 100644 mediapipe/util/frame_buffer/rgb_buffer.cc create mode 100644 mediapipe/util/frame_buffer/rgb_buffer.h create mode 100644 mediapipe/util/frame_buffer/rgb_buffer_test.cc create mode 100644 mediapipe/util/frame_buffer/yuv_buffer.cc create mode 100644 mediapipe/util/frame_buffer/yuv_buffer.h create mode 100644 mediapipe/util/frame_buffer/yuv_buffer_test.cc create mode 100644 third_party/halide.BUILD create mode 100644 third_party/halide/BUILD create mode 100644 third_party/halide/BUILD.bazel create mode 100644 third_party/halide/halide.bzl diff --git a/WORKSPACE b/WORKSPACE index 10f0c1ac5..2f2bc3d9d 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -543,3 +543,35 @@ external_files() load("@//third_party:wasm_files.bzl", "wasm_files") wasm_files() + +# Halide + +new_local_repository( + name = "halide", + build_file = "@//third_party/halide:BUILD.bazel", + path = "third_party/halide" +) + +http_archive( + name = "linux_halide", + sha256 = "be3bdd067acb9ee0d37d0830821113cd69174bee46da466a836d8829fef7cf91", + strip_prefix = "Halide-14.0.0-x86-64-linux", + urls = ["https://github.com/halide/Halide/releases/download/v14.0.0/Halide-14.0.0-x86-64-linux-6b9ed2afd1d6d0badf04986602c943e287d44e46.tar.gz"], + build_file = "@//third_party:halide.BUILD", +) + +http_archive( + name = "macos_halide", + sha256 = "569b1cda858d278d9dd68e072768d8347775e419f2ff39ff34d001ceb299e3c4", + strip_prefix = "Halide-14.0.0-x86-64-osx", + urls = ["https://github.com/halide/Halide/releases/download/v14.0.0/Halide-14.0.0-x86-64-osx-6b9ed2afd1d6d0badf04986602c943e287d44e46.tar.gz"], + build_file = "@//third_party:halide.BUILD", +) + +http_archive( + name = "windows_halide", + sha256 = "a7a0481af2691ec436d79c20ca441e9a701bfce409f4f763dab75a8f1d740179", + strip_prefix = "Halide-14.0.0-x86-64-windows", + urls = ["https://github.com/halide/Halide/releases/download/v14.0.0/Halide-14.0.0-x86-64-windows-6b9ed2afd1d6d0badf04986602c943e287d44e46.zip"], + build_file = "@//third_party:halide.BUILD", +) diff --git a/mediapipe/util/frame_buffer/BUILD b/mediapipe/util/frame_buffer/BUILD new file mode 100644 index 000000000..27343d6df --- /dev/null +++ b/mediapipe/util/frame_buffer/BUILD @@ -0,0 +1,106 @@ +# 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. + +package(default_visibility = ["//mediapipe/util/frame_buffer:__subpackages__"]) + +cc_library( + name = "frame_buffer_util", + srcs = ["frame_buffer_util.cc"], + hdrs = ["frame_buffer_util.h"], + visibility = ["//visibility:public"], + deps = [ + ":buffer", + "//mediapipe/framework/formats:frame_buffer", + "//mediapipe/framework/port:status", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", + ], +) + +cc_test( + name = "frame_buffer_util_test", + srcs = [ + "frame_buffer_util_test.cc", + ], + deps = [ + ":frame_buffer_util", + "//mediapipe/framework/formats:frame_buffer", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:status", + ], +) + +cc_library( + name = "buffer", + srcs = [ + "buffer_common.cc", + "gray_buffer.cc", + "rgb_buffer.cc", + "yuv_buffer.cc", + ], + hdrs = [ + "buffer_common.h", + "gray_buffer.h", + "rgb_buffer.h", + "yuv_buffer.h", + ], + deps = [ + "//mediapipe/util/frame_buffer/halide:gray_flip_halide", + "//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_gray_halide", + "//mediapipe/util/frame_buffer/halide:rgb_resize_halide", + "//mediapipe/util/frame_buffer/halide:rgb_rgb_halide", + "//mediapipe/util/frame_buffer/halide:rgb_rotate_halide", + "//mediapipe/util/frame_buffer/halide:rgb_yuv_halide", + "//mediapipe/util/frame_buffer/halide:yuv_flip_halide", + "//mediapipe/util/frame_buffer/halide:yuv_resize_halide", + "//mediapipe/util/frame_buffer/halide:yuv_rgb_halide", + "//mediapipe/util/frame_buffer/halide:yuv_rotate_halide", + "@halide//:runtime", + ], +) + +# Tests: +cc_test( + name = "rgb_buffer_test", + srcs = ["rgb_buffer_test.cc"], + deps = [ + ":buffer", + "//mediapipe/framework/port:gtest_main", + "@com_google_absl//absl/log", + ], +) + +cc_test( + name = "yuv_buffer_test", + srcs = ["yuv_buffer_test.cc"], + deps = [ + ":buffer", + "//mediapipe/framework/port:gtest_main", + "@com_google_absl//absl/log", + ], +) + +cc_test( + name = "gray_buffer_test", + srcs = ["gray_buffer_test.cc"], + deps = [ + ":buffer", + "//mediapipe/framework/port:gtest_main", + "@com_google_absl//absl/log", + ], +) diff --git a/mediapipe/util/frame_buffer/buffer_common.cc b/mediapipe/util/frame_buffer/buffer_common.cc new file mode 100644 index 000000000..180f52c03 --- /dev/null +++ b/mediapipe/util/frame_buffer/buffer_common.cc @@ -0,0 +1,40 @@ +// 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/buffer_common.h" + +namespace mediapipe { +namespace frame_buffer { +namespace common { + +bool crop_buffer(int x0, int y0, int x1, int y1, halide_buffer_t* buffer) { + if (x0 < 0 || x1 >= buffer->dim[0].extent) { + return false; + } + if (y0 < 0 || y1 >= buffer->dim[1].extent) { + return false; + } + + // Move the start pointer so that it points at (x0, y0) and set the new + // extents. Leave the strides unchanged; we simply skip over the cropped + // image data. + buffer->host += y0 * buffer->dim[1].stride + x0 * buffer->dim[0].stride; + buffer->dim[0].extent = x1 - x0 + 1; + buffer->dim[1].extent = y1 - y0 + 1; + return true; +} + +} // namespace common +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/buffer_common.h b/mediapipe/util/frame_buffer/buffer_common.h new file mode 100644 index 000000000..9e0e891d3 --- /dev/null +++ b/mediapipe/util/frame_buffer/buffer_common.h @@ -0,0 +1,32 @@ +// 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_BUFFER_COMMON_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_BUFFER_COMMON_H_ + +#include "HalideRuntime.h" + +namespace mediapipe { +namespace frame_buffer { +namespace common { + +// Performs in-place cropping on the given buffer; the provided rectangle +// becomes the full extent of the buffer upon success. Returns false on error. +bool crop_buffer(int x0, int y0, int x1, int y1, halide_buffer_t* buffer); + +} // namespace common +} // namespace frame_buffer +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_BUFFER_COMMON_H_ diff --git a/mediapipe/util/frame_buffer/frame_buffer_util.cc b/mediapipe/util/frame_buffer/frame_buffer_util.cc new file mode 100644 index 000000000..b18e8cb13 --- /dev/null +++ b/mediapipe/util/frame_buffer/frame_buffer_util.cc @@ -0,0 +1,794 @@ +// 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/frame_buffer_util.h" + +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/port/status_macros.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" + +namespace mediapipe { +namespace frame_buffer { + +namespace { + +constexpr int kRgbaChannels = 4; +constexpr int kRgbaPixelBytes = 4; +constexpr int kRgbChannels = 3; +constexpr int kRgbPixelBytes = 3; +constexpr int kGrayChannel = 1; +constexpr int kGrayPixelBytes = 1; + +// YUV helpers. +//------------------------------------------------------------------------------ + +// Returns whether the buffer is part of the supported Yuv format. +bool IsSupportedYuvBuffer(const FrameBuffer& buffer) { + return buffer.format() == FrameBuffer::Format::kNV21 || + buffer.format() == FrameBuffer::Format::kNV12 || + buffer.format() == FrameBuffer::Format::kYV12 || + buffer.format() == FrameBuffer::Format::kYV21; +} + +// Shared validation functions. +//------------------------------------------------------------------------------ + +// Indicates whether the given buffers have the same dimensions. +bool AreBufferDimsEqual(const FrameBuffer& buffer1, + const FrameBuffer& buffer2) { + return buffer1.dimension() == buffer2.dimension(); +} + +// Indicates whether the given buffers formats are compatible. Same formats are +// compatible and all YUV family formats (e.g. NV21, NV12, YV12, YV21, etc) are +// compatible. +bool AreBufferFormatsCompatible(const FrameBuffer& buffer1, + const FrameBuffer& buffer2) { + switch (buffer1.format()) { + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return (buffer2.format() == FrameBuffer::Format::kRGBA || + buffer2.format() == FrameBuffer::Format::kRGB); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return (buffer2.format() == FrameBuffer::Format::kNV12 || + buffer2.format() == FrameBuffer::Format::kNV21 || + buffer2.format() == FrameBuffer::Format::kYV12 || + buffer2.format() == FrameBuffer::Format::kYV21); + case FrameBuffer::Format::kGRAY: + default: + return buffer1.format() == buffer2.format(); + } +} + +absl::Status ValidateBufferFormat(const FrameBuffer& buffer) { + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + case FrameBuffer::Format::kRGB: + case FrameBuffer::Format::kRGBA: + if (buffer.plane_count() == 1) return absl::OkStatus(); + return absl::InvalidArgumentError( + "Plane count must be 1 for grayscale and RGB[a] buffers."); + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kYV21: + case FrameBuffer::Format::kYV12: + return absl::OkStatus(); + default: + return absl::InternalError( + absl::StrFormat("Unsupported buffer format: %i.", buffer.format())); + } +} + +absl::Status ValidateBufferFormats(const FrameBuffer& buffer1, + const FrameBuffer& buffer2) { + MP_RETURN_IF_ERROR(ValidateBufferFormat(buffer1)); + MP_RETURN_IF_ERROR(ValidateBufferFormat(buffer2)); + return absl::OkStatus(); +} + +absl::Status ValidateResizeBufferInputs(const FrameBuffer& buffer, + const FrameBuffer& output_buffer) { + bool valid_format = false; + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + case FrameBuffer::Format::kRGB: + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + valid_format = (buffer.format() == output_buffer.format()); + break; + case FrameBuffer::Format::kRGBA: + valid_format = (output_buffer.format() == FrameBuffer::Format::kRGBA || + output_buffer.format() == FrameBuffer::Format::kRGB); + break; + default: + return absl::InternalError( + absl::StrFormat("Unsupported buffer format: %i.", buffer.format())); + } + if (!valid_format) { + return absl::InvalidArgumentError( + "Input and output buffer formats must match."); + } + return ValidateBufferFormats(buffer, output_buffer); +} + +absl::Status ValidateRotateBufferInputs(const FrameBuffer& buffer, + const FrameBuffer& output_buffer, + int angle_deg) { + if (!AreBufferFormatsCompatible(buffer, output_buffer)) { + return absl::InvalidArgumentError( + "Input and output buffer formats must match."); + } + + const bool is_dimension_change = (angle_deg / 90) % 2 == 1; + const bool are_dimensions_rotated = + (buffer.dimension().width == output_buffer.dimension().height) && + (buffer.dimension().height == output_buffer.dimension().width); + const bool are_dimensions_equal = + buffer.dimension() == output_buffer.dimension(); + + if (angle_deg >= 360 || angle_deg <= 0 || angle_deg % 90 != 0) { + return absl::InvalidArgumentError( + "Rotation angle must be between 0 and 360, in multiples of 90 " + "degrees."); + } else if ((is_dimension_change && !are_dimensions_rotated) || + (!is_dimension_change && !are_dimensions_equal)) { + return absl::InvalidArgumentError( + "Output buffer has invalid dimensions for rotation."); + } + return absl::OkStatus(); +} + +absl::Status ValidateCropBufferInputs(const FrameBuffer& buffer, + const FrameBuffer& output_buffer, int x0, + int y0, int x1, int y1) { + if (!AreBufferFormatsCompatible(buffer, output_buffer)) { + return absl::InvalidArgumentError( + "Input and output buffer formats must match."); + } + + bool is_buffer_size_valid = + ((x1 < buffer.dimension().width) && y1 < buffer.dimension().height); + bool are_points_valid = (x0 >= 0) && (y0 >= 0) && (x1 >= x0) && (y1 >= y0); + + if (!is_buffer_size_valid || !are_points_valid) { + return absl::InvalidArgumentError("Invalid crop coordinates."); + } + return absl::OkStatus(); +} + +absl::Status ValidateFlipBufferInputs(const FrameBuffer& buffer, + const FrameBuffer& output_buffer) { + if (!AreBufferFormatsCompatible(buffer, output_buffer)) { + return absl::InvalidArgumentError( + "Input and output buffer formats must match."); + } + return AreBufferDimsEqual(buffer, output_buffer) + ? absl::OkStatus() + : absl::InvalidArgumentError( + "Input and output buffers must have the same dimensions."); +} + +absl::Status ValidateConvertFormats(FrameBuffer::Format from_format, + FrameBuffer::Format to_format) { + if (from_format == to_format) { + return absl::InvalidArgumentError("Formats must be different."); + } + + switch (from_format) { + case FrameBuffer::Format::kGRAY: + return absl::InvalidArgumentError( + "Grayscale format does not convert to other formats."); + case FrameBuffer::Format::kRGB: + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return absl::OkStatus(); + default: + return absl::InternalError( + absl::StrFormat("Unsupported buffer format: %i.", from_format)); + } +} + +// Construct buffer helper functions. +//------------------------------------------------------------------------------ + +// Creates NV12 / NV21 / YV12 / YV21 YuvBuffer from the input `buffer`. The +// output YuvBuffer is agnostic to the YUV format since the YUV buffers are +// managed individually. +absl::StatusOr CreateYuvBuffer(const FrameBuffer& buffer) { + ASSIGN_OR_RETURN(FrameBuffer::YuvData yuv_data, + FrameBuffer::GetYuvDataFromFrameBuffer(buffer)); + return YuvBuffer(const_cast(yuv_data.y_buffer), + const_cast(yuv_data.u_buffer), + const_cast(yuv_data.v_buffer), + buffer.dimension().width, buffer.dimension().height, + yuv_data.y_row_stride, yuv_data.uv_row_stride, + yuv_data.uv_pixel_stride); +} + +absl::StatusOr CreateGrayBuffer(const FrameBuffer& buffer) { + if (buffer.plane_count() != 1) { + return absl::InternalError("Unsupported grayscale planar format."); + } + return GrayBuffer(const_cast(buffer.plane(0).buffer()), + buffer.dimension().width, buffer.dimension().height); +} + +absl::StatusOr CreateRgbBuffer(const FrameBuffer& buffer) { + if (buffer.plane_count() != 1) { + return absl::InternalError("Unsupported rgb[a] planar format."); + } + bool alpha = buffer.format() == FrameBuffer::Format::kRGBA ? true : false; + return RgbBuffer(const_cast(buffer.plane(0).buffer()), + buffer.dimension().width, buffer.dimension().height, + buffer.plane(0).stride().row_stride_bytes, alpha); +} + +// Grayscale transformation functions. +//------------------------------------------------------------------------------ + +absl::Status CropGrayscale(const FrameBuffer& buffer, int x0, int y0, int x1, + int y1, FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer)); + bool success_crop = input.Crop(x0, y0, x1, y1); + if (!success_crop) { + return absl::UnknownError("Halide grayscale crop operation failed."); + } + bool success_resize = input.Resize(&output); + if (!success_resize) { + return absl::UnknownError("Halide grayscale resize operation failed."); + } + return absl::OkStatus(); +} + +absl::Status ResizeGrayscale(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer)); + return input.Resize(&output) + ? absl::OkStatus() + : absl::UnknownError("Halide grayscale resize operation failed."); +} + +absl::Status RotateGrayscale(const FrameBuffer& buffer, int angle_deg, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer)); + return input.Rotate(angle_deg % 360, &output) + ? absl::OkStatus() + : absl::UnknownError("Halide grayscale rotate operation failed."); +} + +absl::Status FlipHorizontallyGrayscale(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer)); + return input.FlipHorizontally(&output) + ? absl::OkStatus() + : absl::UnknownError( + "Halide grayscale horizontal flip operation failed."); +} + +absl::Status FlipVerticallyGrayscale(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateGrayBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer)); + return input.FlipVertically(&output) + ? absl::OkStatus() + : absl::UnknownError( + "Halide grayscale vertical flip operation failed."); +} + +// Rgb transformation functions. +//------------------------------------------------------------------------------ + +absl::Status ResizeRgb(const FrameBuffer& buffer, FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + return input.Resize(&output) + ? absl::OkStatus() + : absl::UnknownError("Halide rgb[a] resize operation failed."); +} + +absl::Status ConvertRgb(const FrameBuffer& buffer, FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + bool result = false; + if (output_buffer->format() == FrameBuffer::Format::kGRAY) { + ASSIGN_OR_RETURN(auto output, CreateGrayBuffer(*output_buffer)); + result = input.Convert(&output); + } else if (IsSupportedYuvBuffer(*output_buffer)) { + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + result = input.Convert(&output); + } else if (output_buffer->format() == FrameBuffer::Format::kRGBA || + output_buffer->format() == FrameBuffer::Format::kRGB) { + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + result = input.Convert(&output); + } + return result ? absl::OkStatus() + : absl::UnknownError("Halide rgb[a] convert operation failed."); +} + +absl::Status CropRgb(const FrameBuffer& buffer, int x0, int y0, int x1, int y1, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + bool success_crop = input.Crop(x0, y0, x1, y1); + if (!success_crop) { + return absl::UnknownError("Halide rgb[a] crop operation failed."); + } + bool success_resize = input.Resize(&output); + if (!success_resize) { + return absl::UnknownError("Halide rgb resize operation failed."); + } + return absl::OkStatus(); +} + +absl::Status FlipHorizontallyRgb(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + return input.FlipHorizontally(&output) + ? absl::OkStatus() + : absl::UnknownError( + "Halide rgb[a] horizontal flip operation failed."); +} + +absl::Status FlipVerticallyRgb(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + return input.FlipVertically(&output) + ? absl::OkStatus() + : absl::UnknownError( + "Halide rgb[a] vertical flip operation failed."); +} + +absl::Status RotateRgb(const FrameBuffer& buffer, int angle, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateRgbBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + return input.Rotate(angle % 360, &output) + ? absl::OkStatus() + : absl::UnknownError("Halide rgb[a] rotate operation failed."); +} + +// Yuv transformation functions. +//------------------------------------------------------------------------------ + +absl::Status CropYuv(const FrameBuffer& buffer, int x0, int y0, int x1, int y1, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + bool success_crop = input.Crop(x0, y0, x1, y1); + if (!success_crop) { + return absl::UnknownError("Halide YUV crop operation failed."); + } + bool success_resize = input.Resize(&output); + if (!success_resize) { + return absl::UnknownError("Halide YUV resize operation failed."); + } + return absl::OkStatus(); +} + +absl::Status ResizeYuv(const FrameBuffer& buffer, FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + return input.Resize(&output) + ? absl::OkStatus() + : absl::UnknownError("Halide YUV resize operation failed."); +} + +absl::Status RotateYuv(const FrameBuffer& buffer, int angle_deg, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + return input.Rotate(angle_deg % 360, &output) + ? absl::OkStatus() + : absl::UnknownError("Halide YUV rotate operation failed."); +} + +absl::Status FlipHorizontallyYuv(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + return input.FlipHorizontally(&output) + ? absl::OkStatus() + : absl::UnknownError( + "Halide YUV horizontal flip operation failed."); +} + +absl::Status FlipVerticallyYuv(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer)); + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + return input.FlipVertically(&output) + ? absl::OkStatus() + : absl::UnknownError("Halide YUV vertical flip operation failed."); +} + +// Converts input YUV `buffer` into the `output_buffer` in RGB, RGBA or gray +// scale format. +absl::Status ConvertYuv(const FrameBuffer& buffer, FrameBuffer* output_buffer) { + bool success_convert = false; + ASSIGN_OR_RETURN(auto input, CreateYuvBuffer(buffer)); + if (output_buffer->format() == FrameBuffer::Format::kRGBA || + output_buffer->format() == FrameBuffer::Format::kRGB) { + ASSIGN_OR_RETURN(auto output, CreateRgbBuffer(*output_buffer)); + bool half_sampling = false; + if (buffer.dimension().width / 2 == output_buffer->dimension().width && + buffer.dimension().height / 2 == output_buffer->dimension().height) { + half_sampling = true; + } + success_convert = input.Convert(half_sampling, &output); + } else if (output_buffer->format() == FrameBuffer::Format::kGRAY) { + if (buffer.plane(0).stride().row_stride_bytes == buffer.dimension().width) { + std::copy(input.y_buffer()->host, + input.y_buffer()->host + buffer.dimension().Size(), + const_cast(output_buffer->plane(0).buffer())); + } else { + // The y_buffer is padded. The conversion removes the padding. + uint8_t* gray_buffer = + const_cast(output_buffer->plane(0).buffer()); + for (int i = 0; i < buffer.dimension().height; i++) { + int src_address = i * buffer.plane(0).stride().row_stride_bytes; + int dest_address = i * buffer.dimension().width; + std::memcpy(&gray_buffer[dest_address], + &buffer.plane(0).buffer()[src_address], + buffer.dimension().width); + } + } + success_convert = true; + } else if (IsSupportedYuvBuffer(*output_buffer)) { + ASSIGN_OR_RETURN(auto output, CreateYuvBuffer(*output_buffer)); + success_convert = input.Resize(&output); + } + return success_convert + ? absl::OkStatus() + : absl::UnknownError("Halide YUV convert operation failed."); +} + +} // namespace + +// Public methods. +//------------------------------------------------------------------------------ + +std::shared_ptr CreateFromRgbaRawBuffer( + uint8_t* input, FrameBuffer::Dimension dimension, + FrameBuffer::Stride stride) { + if (stride == kDefaultStride) { + stride.row_stride_bytes = dimension.width * kRgbaChannels; + stride.pixel_stride_bytes = kRgbaChannels; + } + FrameBuffer::Plane input_plane(/*buffer=*/input, + /*stride=*/stride); + std::vector planes{input_plane}; + return std::make_shared(planes, dimension, + FrameBuffer::Format::kRGBA); +} + +std::shared_ptr CreateFromRgbRawBuffer( + uint8_t* input, FrameBuffer::Dimension dimension, + FrameBuffer::Stride stride) { + if (stride == kDefaultStride) { + stride.row_stride_bytes = dimension.width * kRgbChannels; + stride.pixel_stride_bytes = kRgbChannels; + } + FrameBuffer::Plane input_plane(/*buffer=*/input, + /*stride=*/stride); + std::vector planes{input_plane}; + return std::make_shared(planes, dimension, + FrameBuffer::Format::kRGB); +} + +std::shared_ptr CreateFromGrayRawBuffer( + uint8_t* input, FrameBuffer::Dimension dimension, + FrameBuffer::Stride stride) { + if (stride == kDefaultStride) { + stride.row_stride_bytes = dimension.width * kGrayChannel; + stride.pixel_stride_bytes = kGrayChannel; + } + FrameBuffer::Plane input_plane(/*buffer=*/input, + /*stride=*/stride); + std::vector planes{input_plane}; + return std::make_shared(planes, dimension, + FrameBuffer::Format::kGRAY); +} + +absl::StatusOr> CreateFromYuvRawBuffer( + uint8_t* y_plane, uint8_t* u_plane, uint8_t* v_plane, + FrameBuffer::Format format, FrameBuffer::Dimension dimension, + int row_stride_y, int row_stride_uv, int pixel_stride_uv) { + const int pixel_stride_y = 1; + std::vector planes; + if (format == FrameBuffer::Format::kNV21 || + format == FrameBuffer::Format::kYV12) { + planes = {{y_plane, /*stride=*/{row_stride_y, pixel_stride_y}}, + {v_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}}, + {u_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}}}; + } else if (format == FrameBuffer::Format::kNV12 || + format == FrameBuffer::Format::kYV21) { + planes = {{y_plane, /*stride=*/{row_stride_y, pixel_stride_y}}, + {u_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}}, + {v_plane, /*stride=*/{row_stride_uv, pixel_stride_uv}}}; + } else { + return absl::InvalidArgumentError( + absl::StrFormat("Input format is not YUV-like: %i.", format)); + } + return std::make_shared(planes, dimension, format); +} + +absl::StatusOr> CreateFromRawBuffer( + uint8_t* buffer, FrameBuffer::Dimension dimension, + const FrameBuffer::Format target_format) { + switch (target_format) { + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: { + FrameBuffer::Plane plane(/*buffer=*/buffer, + /*stride=*/{dimension.width, kGrayChannel}); + std::vector planes{plane}; + return std::make_shared(planes, dimension, target_format); + } + case FrameBuffer::Format::kYV12: { + ASSIGN_OR_RETURN(const FrameBuffer::Dimension uv_dimension, + GetUvPlaneDimension(dimension, target_format)); + return CreateFromYuvRawBuffer( + /*y_plane=*/buffer, + /*u_plane=*/buffer + dimension.Size() + uv_dimension.Size(), + /*v_plane=*/buffer + dimension.Size(), target_format, dimension, + /*row_stride_y=*/dimension.width, uv_dimension.width, + /*pixel_stride_uv=*/1); + } + case FrameBuffer::Format::kYV21: { + ASSIGN_OR_RETURN(const FrameBuffer::Dimension uv_dimension, + GetUvPlaneDimension(dimension, target_format)); + return CreateFromYuvRawBuffer( + /*y_plane=*/buffer, /*u_plane=*/buffer + dimension.Size(), + /*v_plane=*/buffer + dimension.Size() + uv_dimension.Size(), + target_format, dimension, /*row_stride_y=*/dimension.width, + uv_dimension.width, + /*pixel_stride_uv=*/1); + } + case FrameBuffer::Format::kRGBA: + return CreateFromRgbaRawBuffer(buffer, dimension); + case FrameBuffer::Format::kRGB: + return CreateFromRgbRawBuffer(buffer, dimension); + case FrameBuffer::Format::kGRAY: + return CreateFromGrayRawBuffer(buffer, dimension); + default: + return absl::InternalError( + absl::StrFormat("Unsupported buffer format: %i.", target_format)); + } +} + +absl::Status Crop(const FrameBuffer& buffer, int x0, int y0, int x1, int y1, + FrameBuffer* output_buffer) { + MP_RETURN_IF_ERROR( + ValidateCropBufferInputs(buffer, *output_buffer, x0, y0, x1, y1)); + MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer)); + + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + return CropGrayscale(buffer, x0, y0, x1, y1, output_buffer); + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return CropRgb(buffer, x0, y0, x1, y1, output_buffer); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return CropYuv(buffer, x0, y0, x1, y1, output_buffer); + default: + return absl::InternalError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + +absl::Status Resize(const FrameBuffer& buffer, FrameBuffer* output_buffer) { + MP_RETURN_IF_ERROR(ValidateResizeBufferInputs(buffer, *output_buffer)); + + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + return ResizeGrayscale(buffer, output_buffer); + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return ResizeRgb(buffer, output_buffer); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return ResizeYuv(buffer, output_buffer); + default: + return absl::InternalError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + +absl::Status Rotate(const FrameBuffer& buffer, int angle_deg, + FrameBuffer* output_buffer) { + MP_RETURN_IF_ERROR( + ValidateRotateBufferInputs(buffer, *output_buffer, angle_deg)); + MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer)); + + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + return RotateGrayscale(buffer, angle_deg, output_buffer); + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return RotateRgb(buffer, angle_deg, output_buffer); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return RotateYuv(buffer, angle_deg, output_buffer); + default: + return absl::InternalError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + +absl::Status FlipHorizontally(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + MP_RETURN_IF_ERROR(ValidateFlipBufferInputs(buffer, *output_buffer)); + MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer)); + + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + return FlipHorizontallyGrayscale(buffer, output_buffer); + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return FlipHorizontallyRgb(buffer, output_buffer); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return FlipHorizontallyYuv(buffer, output_buffer); + default: + return absl::InternalError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + +absl::Status FlipVertically(const FrameBuffer& buffer, + FrameBuffer* output_buffer) { + MP_RETURN_IF_ERROR(ValidateFlipBufferInputs(buffer, *output_buffer)); + MP_RETURN_IF_ERROR(ValidateBufferFormats(buffer, *output_buffer)); + + switch (buffer.format()) { + case FrameBuffer::Format::kGRAY: + return FlipVerticallyGrayscale(buffer, output_buffer); + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return FlipVerticallyRgb(buffer, output_buffer); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return FlipVerticallyYuv(buffer, output_buffer); + default: + return absl::InternalError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + +absl::Status Convert(const FrameBuffer& buffer, FrameBuffer* output_buffer) { + MP_RETURN_IF_ERROR( + ValidateConvertFormats(buffer.format(), output_buffer->format())); + + switch (buffer.format()) { + case FrameBuffer::Format::kRGBA: + case FrameBuffer::Format::kRGB: + return ConvertRgb(buffer, output_buffer); + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return ConvertYuv(buffer, output_buffer); + default: + return absl::InternalError( + absl::StrFormat("Format %i is not supported.", buffer.format())); + } +} + +int GetFrameBufferByteSize(FrameBuffer::Dimension dimension, + FrameBuffer::Format format) { + switch (format) { + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return /*y plane*/ dimension.Size() + + /*uv plane*/ (dimension.width + 1) / 2 * (dimension.height + 1) / + 2 * 2; + case FrameBuffer::Format::kRGB: + return dimension.Size() * kRgbPixelBytes; + case FrameBuffer::Format::kRGBA: + return dimension.Size() * kRgbaPixelBytes; + case FrameBuffer::Format::kGRAY: + return dimension.Size(); + default: + return 0; + } +} + +absl::StatusOr GetPixelStrides(FrameBuffer::Format format) { + switch (format) { + case FrameBuffer::Format::kGRAY: + return kGrayPixelBytes; + case FrameBuffer::Format::kRGB: + return kRgbPixelBytes; + case FrameBuffer::Format::kRGBA: + return kRgbaPixelBytes; + default: + return absl::InvalidArgumentError(absl::StrFormat( + "GetPixelStrides does not support format: %i.", format)); + } +} + +absl::StatusOr GetUvRawBuffer(const FrameBuffer& buffer) { + if (buffer.format() != FrameBuffer::Format::kNV12 && + buffer.format() != FrameBuffer::Format::kNV21) { + return absl::InvalidArgumentError( + "Only support getting biplanar UV buffer from NV12/NV21 frame buffer."); + } + ASSIGN_OR_RETURN(FrameBuffer::YuvData yuv_data, + FrameBuffer::GetYuvDataFromFrameBuffer(buffer)); + const uint8_t* uv_buffer = buffer.format() == FrameBuffer::Format::kNV12 + ? yuv_data.u_buffer + : yuv_data.v_buffer; + return uv_buffer; +} + +absl::StatusOr GetUvPlaneDimension( + FrameBuffer::Dimension dimension, FrameBuffer::Format format) { + if (dimension.width <= 0 || dimension.height <= 0) { + return absl::InvalidArgumentError( + absl::StrFormat("Invalid input dimension: {%d, %d}.", dimension.width, + dimension.height)); + } + switch (format) { + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: + return FrameBuffer::Dimension{(dimension.width + 1) / 2, + (dimension.height + 1) / 2}; + default: + return absl::InvalidArgumentError( + absl::StrFormat("Input format is not YUV-like: %i.", format)); + } +} + +FrameBuffer::Dimension GetCropDimension(int x0, int x1, int y0, int y1) { + return {x1 - x0 + 1, y1 - y0 + 1}; +} + +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/frame_buffer_util.h b/mediapipe/util/frame_buffer/frame_buffer_util.h new file mode 100644 index 000000000..eb17e0094 --- /dev/null +++ b/mediapipe/util/frame_buffer/frame_buffer_util.h @@ -0,0 +1,131 @@ +// 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_FRAME_BUFFER_UTIL_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_FRAME_BUFFER_UTIL_H_ + +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "mediapipe/framework/formats/frame_buffer.h" + +namespace mediapipe { +namespace frame_buffer { + +// Creation helpers. +//------------------------------------------------------------------------------ + +// Default stride value for creating frame buffer from raw buffer. When using +// this default value, the default row stride and pixel stride values will be +// applied. e.g. for an RGB image: +// row_stride = width * 3 +// pixel_stride = 3. +inline constexpr FrameBuffer::Stride kDefaultStride = {0, 0}; + +// Creates a FrameBuffer from raw RGBA buffer and passing arguments. +std::shared_ptr CreateFromRgbaRawBuffer( + uint8_t* input, FrameBuffer::Dimension dimension, + FrameBuffer::Stride stride = kDefaultStride); + +// Creates a FrameBuffer from raw RGB buffer and passing arguments. +std::shared_ptr CreateFromRgbRawBuffer( + uint8_t* input, FrameBuffer::Dimension dimension, + FrameBuffer::Stride stride = kDefaultStride); + +// Creates a FrameBuffer from raw grayscale buffer and passing arguments. +std::shared_ptr CreateFromGrayRawBuffer( + uint8_t* input, FrameBuffer::Dimension dimension, + FrameBuffer::Stride stride = kDefaultStride); + +// Creates a FrameBuffer from raw YUV buffer and passing arguments. +absl::StatusOr> CreateFromYuvRawBuffer( + uint8_t* y_plane, uint8_t* u_plane, uint8_t* v_plane, + FrameBuffer::Format format, FrameBuffer::Dimension dimension, + int row_stride_y, int row_stride_uv, int pixel_stride_uv); + +// Creates an instance of FrameBuffer from raw buffer and passing arguments. +absl::StatusOr> CreateFromRawBuffer( + uint8_t* buffer, FrameBuffer::Dimension dimension, + FrameBuffer::Format target_format); + +// Transformations. +//------------------------------------------------------------------------------ + +// Crops `buffer` to the specified points. +// +// (x0, y0) represents the top-left point of the buffer. +// (x1, y1) represents the bottom-right point of the buffer. +// +// The implementation performs origin moving and resizing operations. +absl::Status Crop(const FrameBuffer& buffer, int x0, int y0, int x1, int y1, + FrameBuffer* output_buffer); + +// Resizes `buffer` to the size of the given `output_buffer` using bilinear +// interpolation. +absl::Status Resize(const FrameBuffer& buffer, FrameBuffer* output_buffer); + +// Rotates `buffer` counter-clockwise by the given `angle_deg` (in degrees). +// +// The given angle must be a multiple of 90 degrees. +absl::Status Rotate(const FrameBuffer& buffer, int angle_deg, + FrameBuffer* output_buffer); + +// Flips `buffer` horizontally. +absl::Status FlipHorizontally(const FrameBuffer& buffer, + FrameBuffer* output_buffer); + +// Flips `buffer` vertically. +absl::Status FlipVertically(const FrameBuffer& buffer, + FrameBuffer* output_buffer); + +// Converts `buffer`'s format to the format of the given `output_buffer`. +// +// Note that grayscale format does not convert to other formats. +// Note the NV21 to RGB/RGBA conversion may downsample by factor of 2 based +// on the buffer and output_buffer dimensions. +absl::Status Convert(const FrameBuffer& buffer, FrameBuffer* output_buffer); + +// Miscellaneous Methods +// ----------------------------------------------------------------- + +// Returns the frame buffer size in bytes based on the input format and +// dimensions. GRAY, YV12/YV21 are in the planar formats, NV12/NV21 are in the +// semi-planar formats with the interleaved UV planes. RGB/RGBA are in the +// interleaved format. +int GetFrameBufferByteSize(FrameBuffer::Dimension dimension, + FrameBuffer::Format format); + +// Returns pixel stride info for kGRAY, kRGB, kRGBA formats. +absl::StatusOr GetPixelStrides(FrameBuffer::Format format); + +// Returns the biplanar UV raw buffer for NV12/NV21 frame buffer. +absl::StatusOr GetUvRawBuffer(const FrameBuffer& buffer); + +// Returns U or V plane dimension with the given buffer `dimension` and +// `format`. Only supports NV12/NV21/YV12/YV21 formats. Returns +// InvalidArgumentError if 'dimension' is invalid or 'format' is other than the +// supported formats. This method assums the UV plane share the same dimension, +// especially for the YV12 / YV21 formats. +absl::StatusOr GetUvPlaneDimension( + FrameBuffer::Dimension dimension, FrameBuffer::Format format); + +// Returns crop dimension based on crop start and end points. +FrameBuffer::Dimension GetCropDimension(int x0, int x1, int y0, int y1); + +} // namespace frame_buffer +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_FRAME_BUFFER_UTIL_H_ diff --git a/mediapipe/util/frame_buffer/frame_buffer_util_test.cc b/mediapipe/util/frame_buffer/frame_buffer_util_test.cc new file mode 100644 index 000000000..d92eb3f53 --- /dev/null +++ b/mediapipe/util/frame_buffer/frame_buffer_util_test.cc @@ -0,0 +1,1176 @@ +// 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/frame_buffer_util.h" + +#include +#include +#include + +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/status_macros.h" +#include "mediapipe/framework/port/status_matchers.h" + +namespace mediapipe { +namespace frame_buffer { +namespace { + +// Grayscale unit tests. +//------------------------------------------------------------------------------ + +TEST(FrameBufferUtil, GrayCrop) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kOutputDimension = {.width = 1, .height = 1}; + uint8_t data[6] = {1, 2, 3, 4, 5, 6}; + uint8_t output_data[2]; + auto input = CreateFromGrayRawBuffer(data, kBufferDimension); + auto output = CreateFromGrayRawBuffer(output_data, kOutputDimension); + + MP_ASSERT_OK(Crop(*input, 0, 1, 0, 1, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 4); +} + +TEST(FrameBufferUtil, GrayResize) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 2, .height = 2}, + kOutputDimension = {.width = 3, .height = 2}; + uint8_t data[4] = {1, 2, 3, 4}; + uint8_t output_data[6]; + auto input = CreateFromGrayRawBuffer(data, kBufferDimension); + auto output = CreateFromGrayRawBuffer(output_data, kOutputDimension); + + MP_ASSERT_OK(Resize(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 1); + EXPECT_EQ(output->plane(0).buffer()[1], 2); + EXPECT_EQ(output->plane(0).buffer()[2], 2); + EXPECT_EQ(output->plane(0).buffer()[3], 3); + EXPECT_EQ(output->plane(0).buffer()[4], 4); + EXPECT_EQ(output->plane(0).buffer()[5], 4); +} + +TEST(FrameBufferUtil, GrayRotate) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kOutputDimension = {.width = 2, .height = 3}; + uint8_t data[6] = {1, 2, 3, 4, 5, 6}; + uint8_t output_data[6]; + auto input = CreateFromGrayRawBuffer(data, kBufferDimension); + auto output = CreateFromGrayRawBuffer(output_data, kOutputDimension); + + MP_ASSERT_OK(Rotate(*input, 90, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 3); + EXPECT_EQ(output->plane(0).buffer()[1], 6); + EXPECT_EQ(output->plane(0).buffer()[2], 2); + EXPECT_EQ(output->plane(0).buffer()[3], 5); + EXPECT_EQ(output->plane(0).buffer()[4], 1); + EXPECT_EQ(output->plane(0).buffer()[5], 4); +} + +TEST(FrameBufferUtil, GrayFlipHorizontally) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}; + uint8_t data[6] = {1, 2, 3, 4, 5, 6}; + uint8_t output_data[6]; + auto input = CreateFromGrayRawBuffer(data, kBufferDimension); + auto output = CreateFromGrayRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(FlipHorizontally(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 3); + EXPECT_EQ(output->plane(0).buffer()[1], 2); + EXPECT_EQ(output->plane(0).buffer()[2], 1); + EXPECT_EQ(output->plane(0).buffer()[3], 6); + EXPECT_EQ(output->plane(0).buffer()[4], 5); + EXPECT_EQ(output->plane(0).buffer()[5], 4); +} + +TEST(FrameBufferUtil, GrayFlipVertically) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}; + uint8_t data[6] = {1, 2, 3, 4, 5, 6}; + uint8_t output_data[6]; + auto input = CreateFromGrayRawBuffer(data, kBufferDimension); + auto output = CreateFromGrayRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(FlipVertically(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 4); + EXPECT_EQ(output->plane(0).buffer()[1], 5); + EXPECT_EQ(output->plane(0).buffer()[2], 6); + EXPECT_EQ(output->plane(0).buffer()[3], 1); + EXPECT_EQ(output->plane(0).buffer()[4], 2); + EXPECT_EQ(output->plane(0).buffer()[5], 3); +} + +// Grayscale EndToEnd tests. +//------------------------------------------------------------------------------ + +struct GrayInputTestParam { + FrameBuffer::Dimension input_dimension; + FrameBuffer::Format input_format; + FrameBuffer::Dimension output_dimension; + FrameBuffer::Format output_format; + int rotation_angle; + int x0; + int y0; + int x1; + int y1; +}; + +enum Operation { + kRotate = 1, + kCrop = 2, + kResize = 3, + kHorizontalFlip = 4, + kVerticalFlip = 5, + kConvert = 6 +}; + +class GrayInputTest : public ::testing::TestWithParam< + std::tuple> {}; + +TEST_P(GrayInputTest, ValidateInputs) { + GrayInputTestParam inputs; + bool is_valid; + Operation operation; + std::tie(operation, inputs, is_valid) = GetParam(); + MP_ASSERT_OK_AND_ASSIGN( + auto input, + CreateFromRawBuffer(/*buffer=*/nullptr, inputs.input_dimension, + inputs.input_format)); + MP_ASSERT_OK_AND_ASSIGN(auto output, + CreateFromRawBuffer(nullptr, inputs.output_dimension, + inputs.output_format)); + + absl::Status status; + switch (operation) { + case kRotate: + status = Rotate(*input, inputs.rotation_angle, output.get()); + break; + case kResize: + status = Resize(*input, output.get()); + break; + case kCrop: { + status = Crop(*input, inputs.x0, inputs.y0, inputs.x1, inputs.y1, + output.get()); + break; + } + case kHorizontalFlip: + status = FlipHorizontally(*input, output.get()); + break; + case kVerticalFlip: + status = FlipVertically(*input, output.get()); + break; + case kConvert: + status = Convert(*input, output.get()); + break; + } + + if (is_valid) { + MP_EXPECT_OK(status); + } else { + EXPECT_THAT(status, StatusIs(absl::StatusCode::kInvalidArgument)); + } +} + +std::tuple CreateGrayRotateInputTestParam( + int in_width, int in_height, FrameBuffer::Format in_format, int out_width, + int out_height, FrameBuffer::Format out_format, int angle, bool is_valid) { + GrayInputTestParam param = { + .input_dimension = FrameBuffer::Dimension{in_width, in_height}, + .input_format = in_format, + .output_dimension = FrameBuffer::Dimension{out_width, out_height}, + .output_format = out_format, + .rotation_angle = angle}; + return std::make_tuple(kRotate, param, is_valid); +} + +INSTANTIATE_TEST_SUITE_P( + ValidateRotateInputs, GrayInputTest, + testing::Values( + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 3, + FrameBuffer::Format::kGRAY, 30, false), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kRGB, 180, false), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, 90, false), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, 0, false), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 3, + FrameBuffer::Format::kGRAY, -90, false), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 3, + FrameBuffer::Format::kGRAY, 90, true), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, 180, true), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 3, + FrameBuffer::Format::kGRAY, 270, true), + CreateGrayRotateInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 3, + FrameBuffer::Format::kGRAY, 450, + false))); + +std::tuple CreateGrayCropInputTestParam( + int in_width, int in_height, FrameBuffer::Format in_format, int out_width, + int out_height, FrameBuffer::Format out_format, int x0, int y0, int x1, + int y1, bool is_valid) { + GrayInputTestParam param = { + .input_dimension = FrameBuffer::Dimension{in_width, in_height}, + .input_format = in_format, + .output_dimension = FrameBuffer::Dimension{out_width, out_height}, + .output_format = out_format, + .x0 = x0, + .y0 = y0, + .x1 = x1, + .y1 = y1}; + return std::make_tuple(kCrop, param, is_valid); +} + +INSTANTIATE_TEST_SUITE_P( + ValidateCropInputs, GrayInputTest, + ::testing::Values( + CreateGrayCropInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kRGB, 0, 0, 3, 2, + false), + CreateGrayCropInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, 1, 1, 1, 4, + false), + CreateGrayCropInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 1, + FrameBuffer::Format::kGRAY, -1, 0, 1, 1, + false), + CreateGrayCropInputTestParam(5, 5, FrameBuffer::Format::kGRAY, 3, 3, + FrameBuffer::Format::kGRAY, 0, 0, 2, 2, + true), + CreateGrayCropInputTestParam(5, 5, FrameBuffer::Format::kGRAY, 2, 2, + FrameBuffer::Format::kGRAY, 1, 2, 2, 3, + true), + CreateGrayCropInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 1, 1, + FrameBuffer::Format::kGRAY, 0, 0, 0, 0, + true))); + +std::tuple CreateGrayResizeInputTestParam( + int in_width, int in_height, FrameBuffer::Format in_format, int out_width, + int out_height, FrameBuffer::Format out_format, bool is_valid) { + GrayInputTestParam param = { + .input_dimension = FrameBuffer::Dimension{in_width, in_height}, + .input_format = in_format, + .output_dimension = FrameBuffer::Dimension{out_width, out_height}, + .output_format = out_format}; + return std::make_tuple(kResize, param, is_valid); +} + +INSTANTIATE_TEST_SUITE_P( + ValidateResizeInputs, GrayInputTest, + ::testing::Values( + CreateGrayResizeInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 1, 1, + FrameBuffer::Format::kRGB, false), + CreateGrayResizeInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 5, 5, + FrameBuffer::Format::kRGB, false), + CreateGrayResizeInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 2, 1, + FrameBuffer::Format::kGRAY, true), + CreateGrayResizeInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 7, 9, + FrameBuffer::Format::kGRAY, true))); + +std::tuple CreateGrayFlipInputTestParam( + int in_width, int in_height, FrameBuffer::Format in_format, int out_width, + int out_height, FrameBuffer::Format out_format, bool horizontal_flip, + bool is_valid) { + GrayInputTestParam param = { + .input_dimension = FrameBuffer::Dimension{in_width, in_height}, + .input_format = in_format, + .output_dimension = FrameBuffer::Dimension{out_width, out_height}, + .output_format = out_format}; + return std::make_tuple(horizontal_flip ? kHorizontalFlip : kVerticalFlip, + param, is_valid); +} + +INSTANTIATE_TEST_SUITE_P( + ValidateFlipInputs, GrayInputTest, + ::testing::Values( + CreateGrayFlipInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kRGB, true, false), + CreateGrayFlipInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 3, + FrameBuffer::Format::kGRAY, true, false), + CreateGrayFlipInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, true, true), + CreateGrayFlipInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kRGB, false, false), + CreateGrayFlipInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 3, + FrameBuffer::Format::kGRAY, false, false), + CreateGrayFlipInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, false, true))); + +std::tuple CreateGrayConvertInputTestParam( + int in_width, int in_height, FrameBuffer::Format in_format, int out_width, + int out_height, FrameBuffer::Format out_format, bool is_valid) { + GrayInputTestParam param = { + .input_dimension = FrameBuffer::Dimension{in_width, in_height}, + .input_format = in_format, + .output_dimension = FrameBuffer::Dimension{out_width, out_height}, + .output_format = out_format}; + return std::make_tuple(kConvert, param, is_valid); +} + +INSTANTIATE_TEST_SUITE_P( + ValidateConvertInputs, GrayInputTest, + ::testing::Values( + CreateGrayConvertInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kRGB, false), + CreateGrayConvertInputTestParam(3, 2, FrameBuffer::Format::kGRAY, 3, 2, + FrameBuffer::Format::kGRAY, false))); + +// Rgb unit tests. +//------------------------------------------------------------------------------ + +struct FrameBufferPlanarFormat { + FrameBufferPlanarFormat() + : format(FrameBuffer::Format::kGRAY), plane_count(0) {} + FrameBufferPlanarFormat(FrameBuffer::Format format, int plane_count) + : format(format), plane_count(plane_count) {} + + FrameBuffer::Format format; + int plane_count; +}; + +class RgbaConvertTest + : public testing::TestWithParam< + std::tuple> { + public: + void SetUp() override { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 2, + .height = 1}; + constexpr int kBufferSize = 20; + std::tie(input_format_, output_planar_format_) = GetParam(); + + // Setup input frame buffer + input_data_ = std::make_unique(kBufferSize); + FrameBuffer::Stride input_stride; + if (input_format_ == FrameBuffer::Format::kRGBA) { + uint8_t data[] = {200, 100, 0, 1, 0, 200, 100, 50}; + std::copy(data, data + 8, input_data_.get()); + input_stride = {/*row_stride_bytes=*/8, /*pixel_stride_bytes=*/4}; + } else { + uint8_t data[] = {200, 100, 0, 0, 200, 100}; + std::copy(data, data + 6, input_data_.get()); + input_stride = {/*row_stride_bytes=*/6, /*pixel_stride_bytes=*/3}; + } + FrameBuffer::Plane input_plane(/*buffer=*/input_data_.get(), input_stride); + std::vector input_planes = {input_plane}; + input_frame_buffer_ = std::make_shared( + input_planes, kBufferDimension, input_format_); + + // Setup output frame buffer + if (output_planar_format_.format == FrameBuffer::Format::kRGBA) { + output_data_1_ = std::make_unique(kBufferSize); + FrameBuffer::Plane output_plane_1( + /*buffer=*/output_data_1_.get(), + /*stride=*/{/*row_stride_bytes=*/8, /*pixel_stride_bytes=*/4}); + std::vector output_planes = {output_plane_1}; + output_frame_buffer_ = std::make_shared( + output_planes, kBufferDimension, output_planar_format_.format); + } else if (output_planar_format_.format == FrameBuffer::Format::kRGB) { + output_data_1_ = std::make_unique(kBufferSize); + FrameBuffer::Plane output_plane_1( + /*buffer=*/output_data_1_.get(), + /*stride=*/{/*row_stride_bytes=*/6, /*pixel_stride_bytes=*/3}); + std::vector output_planes = {output_plane_1}; + output_frame_buffer_ = std::make_shared( + output_planes, kBufferDimension, output_planar_format_.format); + } else if (output_planar_format_.plane_count == 1) { + output_data_1_ = std::make_unique(kBufferSize); + FrameBuffer::Plane output_plane_1( + /*buffer=*/output_data_1_.get(), + /*stride=*/{/*row_stride_bytes=*/2, + /*pixel_stride_bytes=*/1}); + std::vector output_planes = {output_plane_1}; + output_frame_buffer_ = std::make_shared( + output_planes, kBufferDimension, output_planar_format_.format); + } else if (output_planar_format_.plane_count == 2) { + output_data_1_ = std::make_unique(kBufferSize); + FrameBuffer::Plane output_plane_1( + /*buffer=*/output_data_1_.get(), + /*stride=*/{/*row_stride_bytes=*/2, /*pixel_stride_bytes=*/1}); + output_data_2_ = std::make_unique(kBufferSize); + FrameBuffer::Plane output_plane_2( + /*buffer=*/output_data_2_.get(), + /*stride=*/{/*row_stride_bytes=*/1, /*pixel_stride_bytes=*/2}); + std::vector planes = {output_plane_1, output_plane_2}; + output_frame_buffer_ = std::make_shared( + planes, kBufferDimension, output_planar_format_.format); + } else { + output_data_1_ = std::make_unique(kBufferSize); + output_data_2_ = std::make_unique(kBufferSize); + output_data_3_ = std::make_unique(kBufferSize); + FrameBuffer::Plane output_plane_1( + /*buffer=*/output_data_1_.get(), + /*stride=*/{/*row_stride_bytes=*/2, /*pixel_stride_bytes=*/1}); + FrameBuffer::Plane output_plane_2( + /*buffer=*/output_data_2_.get(), + /*stride=*/{/*row_stride_bytes=*/1, /*pixel_stride_bytes=*/1}); + FrameBuffer::Plane output_plane_3( + /*buffer=*/output_data_3_.get(), + /*stride=*/{/*row_stride_bytes=*/1, /*pixel_stride_bytes=*/1}); + std::vector planes = {output_plane_1, output_plane_2, + output_plane_3}; + output_frame_buffer_ = std::make_shared( + planes, kBufferDimension, output_planar_format_.format); + } + } + + protected: + FrameBuffer::Format input_format_; + FrameBufferPlanarFormat output_planar_format_; + + std::unique_ptr output_data_1_; + std::unique_ptr output_data_2_; + std::unique_ptr output_data_3_; + std::unique_ptr input_data_; + + std::shared_ptr input_frame_buffer_; + std::shared_ptr output_frame_buffer_; +}; + +TEST_P(RgbaConvertTest, RgbaToOtherFormatConversion) { + absl::Status status = + Convert(*input_frame_buffer_, output_frame_buffer_.get()); + if (output_planar_format_.format == FrameBuffer::Format::kGRAY) { + MP_ASSERT_OK(status); + EXPECT_EQ(output_data_1_[0], 118); + EXPECT_EQ(output_data_1_[1], 129); + } else if (output_frame_buffer_->format() == FrameBuffer::Format::kNV12 || + output_frame_buffer_->format() == FrameBuffer::Format::kNV21 || + output_frame_buffer_->format() == FrameBuffer::Format::kYV12 || + output_frame_buffer_->format() == FrameBuffer::Format::kYV21) { + MP_ASSERT_OK(status); + MP_ASSERT_OK_AND_ASSIGN( + FrameBuffer::YuvData yuv_data, + FrameBuffer::GetYuvDataFromFrameBuffer(*output_frame_buffer_)); + EXPECT_EQ(yuv_data.y_buffer[0], 118); + EXPECT_EQ(yuv_data.y_buffer[1], 129); + EXPECT_EQ(yuv_data.u_buffer[0], 61); + EXPECT_EQ(yuv_data.v_buffer[0], 186); + } else if (input_format_ == FrameBuffer::Format::kRGBA && + output_frame_buffer_->format() == FrameBuffer::Format::kRGB) { + EXPECT_EQ(output_data_1_[0], 200); + EXPECT_EQ(output_data_1_[1], 100); + EXPECT_EQ(output_data_1_[2], 0); + EXPECT_EQ(output_data_1_[3], 0); + MP_ASSERT_OK(status); + } else if (input_format_ == FrameBuffer::Format::kRGB && + output_frame_buffer_->format() == FrameBuffer::Format::kRGBA) { + MP_ASSERT_OK(status); + EXPECT_EQ(output_data_1_[0], 200); + EXPECT_EQ(output_data_1_[1], 100); + EXPECT_EQ(output_data_1_[2], 0); + EXPECT_EQ(output_data_1_[3], 255); + } else { + ASSERT_FALSE(status.ok()); + } +} + +INSTANTIATE_TEST_SUITE_P( + RgbaToOtherFormatConversion, RgbaConvertTest, + testing::Combine( + testing::Values(FrameBuffer::Format::kRGBA, FrameBuffer::Format::kRGB), + testing::Values(FrameBufferPlanarFormat(FrameBuffer::Format::kGRAY, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kRGBA, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kRGB, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kNV21, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kNV21, + /*plane_count=*/2), + FrameBufferPlanarFormat(FrameBuffer::Format::kNV21, + /*plane_count=*/3), + FrameBufferPlanarFormat(FrameBuffer::Format::kNV12, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kNV12, + /*plane_count=*/2), + FrameBufferPlanarFormat(FrameBuffer::Format::kNV12, + /*plane_count=*/3), + FrameBufferPlanarFormat(FrameBuffer::Format::kYV21, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kYV21, + /*plane_count=*/3), + FrameBufferPlanarFormat(FrameBuffer::Format::kYV12, + /*plane_count=*/1), + FrameBufferPlanarFormat(FrameBuffer::Format::kYV12, + /*plane_count=*/3)))); + +TEST(FrameBufferUtil, RgbaToRgbConversion) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 2, .height = 1}; + uint8_t data[] = {200, 100, 0, 1, 0, 200, 100, 50}; + auto input = CreateFromRgbaRawBuffer(data, kBufferDimension); + uint8_t output_data[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 0}; + auto output = CreateFromRgbRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(Convert(*input, output.get())); + EXPECT_EQ(output_data[0], 200); + EXPECT_EQ(output_data[1], 100); + EXPECT_EQ(output_data[2], 0); + EXPECT_EQ(output_data[3], 0); + EXPECT_EQ(output_data[4], 200); + EXPECT_EQ(output_data[5], 100); +} + +TEST(FrameBufferUtil, RgbaCrop) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kOutputDimension = {.width = 1, .height = 1}; + uint8_t kRgbaTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}; + uint8_t output_data[4]; + auto input = CreateFromRgbaRawBuffer(kRgbaTestData, kBufferDimension); + auto output = CreateFromRgbaRawBuffer(output_data, kOutputDimension); + + MP_ASSERT_OK(Crop(*input, 0, 1, 0, 1, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 13); + EXPECT_EQ(output->plane(0).buffer()[1], 14); + EXPECT_EQ(output->plane(0).buffer()[2], 15); + EXPECT_EQ(output->plane(0).buffer()[3], 16); +} + +TEST(FrameBufferUtil, RgbCrop) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kOutputDimension = {.width = 1, .height = 1}; + uint8_t kRgbTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + uint8_t output_data[3]; + auto input = CreateFromRgbRawBuffer(kRgbTestData, kBufferDimension); + auto output = CreateFromRgbRawBuffer(output_data, kOutputDimension); + + MP_ASSERT_OK(Crop(*input, 0, 1, 0, 1, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 10); + EXPECT_EQ(output->plane(0).buffer()[1], 11); + EXPECT_EQ(output->plane(0).buffer()[2], 12); +} + +TEST(FrameBufferUtil, RgbaFlipHorizontally) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 1}; + uint8_t kRgbaTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}; + uint8_t output_data[sizeof(kRgbaTestData) / 2]; + auto input = CreateFromRgbaRawBuffer(kRgbaTestData, kBufferDimension); + auto output = CreateFromRgbaRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(FlipHorizontally(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 9); + EXPECT_EQ(output->plane(0).buffer()[1], 10); + EXPECT_EQ(output->plane(0).buffer()[2], 11); + EXPECT_EQ(output->plane(0).buffer()[3], 12); + EXPECT_EQ(output->plane(0).buffer()[4], 5); + EXPECT_EQ(output->plane(0).buffer()[5], 6); + EXPECT_EQ(output->plane(0).buffer()[6], 7); + EXPECT_EQ(output->plane(0).buffer()[7], 8); + EXPECT_EQ(output->plane(0).buffer()[8], 1); + EXPECT_EQ(output->plane(0).buffer()[9], 2); + EXPECT_EQ(output->plane(0).buffer()[10], 3); + EXPECT_EQ(output->plane(0).buffer()[11], 4); +} + +TEST(FrameBufferUtil, RgbFlipHorizontally) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 1}; + uint8_t kRgbTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + uint8_t output_data[sizeof(kRgbTestData) / 2]; + auto input = CreateFromRgbRawBuffer(kRgbTestData, kBufferDimension); + auto output = CreateFromRgbRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(FlipHorizontally(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 7); + EXPECT_EQ(output->plane(0).buffer()[1], 8); + EXPECT_EQ(output->plane(0).buffer()[2], 9); + EXPECT_EQ(output->plane(0).buffer()[3], 4); + EXPECT_EQ(output->plane(0).buffer()[4], 5); + EXPECT_EQ(output->plane(0).buffer()[5], 6); + EXPECT_EQ(output->plane(0).buffer()[6], 1); + EXPECT_EQ(output->plane(0).buffer()[7], 2); + EXPECT_EQ(output->plane(0).buffer()[8], 3); +} + +TEST(FrameBufferUtil, RgbaFlipVertically) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}; + uint8_t kRgbaTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}; + uint8_t output_data[sizeof(kRgbaTestData)]; + auto input = CreateFromRgbaRawBuffer(kRgbaTestData, kBufferDimension); + auto output = CreateFromRgbaRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(FlipVertically(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 13); + EXPECT_EQ(output->plane(0).buffer()[1], 14); + EXPECT_EQ(output->plane(0).buffer()[2], 15); + EXPECT_EQ(output->plane(0).buffer()[3], 16); + EXPECT_EQ(output->plane(0).buffer()[12], 1); + EXPECT_EQ(output->plane(0).buffer()[13], 2); + EXPECT_EQ(output->plane(0).buffer()[14], 3); + EXPECT_EQ(output->plane(0).buffer()[15], 4); +} + +TEST(FrameBufferUtil, RgbFlipVertically) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}; + uint8_t kRgbTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + uint8_t output_data[sizeof(kRgbTestData)]; + auto input = CreateFromRgbRawBuffer(kRgbTestData, kBufferDimension); + auto output = CreateFromRgbRawBuffer(output_data, kBufferDimension); + + MP_ASSERT_OK(FlipVertically(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 10); + EXPECT_EQ(output->plane(0).buffer()[1], 11); + EXPECT_EQ(output->plane(0).buffer()[2], 12); + EXPECT_EQ(output->plane(0).buffer()[9], 1); + EXPECT_EQ(output->plane(0).buffer()[10], 2); + EXPECT_EQ(output->plane(0).buffer()[11], 3); +} + +TEST(FrameBufferUtil, RgbaResize) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kResizeUpDimension = {.width = 4, + .height = 2}, + kResizeDownDimension = {.width = 2, + .height = 2}; + uint8_t kRgbaTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}; + uint8_t output_data_up[32]; + auto input = CreateFromRgbaRawBuffer(kRgbaTestData, kBufferDimension); + auto output = CreateFromRgbaRawBuffer(output_data_up, kResizeUpDimension); + + // Test increasing the size. + MP_ASSERT_OK(Resize(*input, output.get())); + uint8_t resize_result_size_increase[] = { + 1, 2, 3, 4, 4, 5, 6, 7, 7, 8, 9, 10, 9, 10, 11, 12, + 13, 14, 15, 16, 16, 17, 18, 19, 19, 20, 21, 22, 21, 22, 23, 24}; + for (int i = 0; i < sizeof(output_data_up); i++) { + EXPECT_EQ(output->plane(0).buffer()[i], resize_result_size_increase[i]); + } + + // Test shrinking the image by half. + uint8_t output_data_down[16]; + output = CreateFromRgbaRawBuffer(output_data_down, kResizeDownDimension); + + MP_ASSERT_OK(Resize(*input, output.get())); + uint8_t resize_result_size_decrease[] = {1, 2, 3, 4, 7, 8, 9, 10, + 13, 14, 15, 16, 19, 20, 21, 22}; + for (int i = 0; i < sizeof(output_data_down); i++) { + EXPECT_EQ(output->plane(0).buffer()[i], resize_result_size_decrease[i]); + } +} + +TEST(FrameBufferUtil, RgbResize) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kResizeUpDimension = {.width = 4, + .height = 3}, + kResizeDownDimension = {.width = 2, + .height = 2}; + uint8_t kRgbTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + auto input = CreateFromRgbRawBuffer(kRgbTestData, kBufferDimension); + + // Test increasing the size. + uint8_t output_data_up[36]; + auto output = CreateFromRgbRawBuffer(output_data_up, kResizeUpDimension); + MP_ASSERT_OK(Resize(*input, output.get())); + uint8_t resize_result_size_increase[] = { + 1, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8, 9, 7, 8, 9, 9, 10, 11, + 11, 12, 13, 13, 14, 15, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16, 17, 18}; + for (int i = 0; i < sizeof(output_data_up); i++) { + EXPECT_EQ(output_data_up[i], resize_result_size_increase[i]); + } + + // Test decreasing the size. + uint8_t output_data_down[12]; + output = CreateFromRgbRawBuffer(output_data_down, kResizeDownDimension); + MP_ASSERT_OK(Resize(*input, output.get())); + + uint8_t resize_result_size_decrease[] = {1, 2, 3, 5, 6, 7, + 10, 11, 12, 14, 15, 16}; + for (int i = 0; i < sizeof(resize_result_size_decrease); i++) { + EXPECT_EQ(output->plane(0).buffer()[i], resize_result_size_decrease[i]); + } +} + +TEST(FrameBufferUtil, RgbaRotate) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kRotatedDimension = {.width = 2, + .height = 3}; + uint8_t kRgbaTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}; + uint8_t output_data[sizeof(kRgbaTestData)]; + auto input = CreateFromRgbaRawBuffer(kRgbaTestData, kBufferDimension); + const std::array kAnglesToTest = {90, 180, 270}; + std::map> kOutputBuffers; + kOutputBuffers[90] = CreateFromRgbaRawBuffer(output_data, kRotatedDimension); + kOutputBuffers[180] = CreateFromRgbaRawBuffer(output_data, kBufferDimension); + kOutputBuffers[270] = CreateFromRgbaRawBuffer(output_data, kRotatedDimension); + const std::map> kRotationResults{ + {90, {9, 10, 11, 12, 21, 22, 23, 24, 5, 6, 7, 8, + 17, 18, 19, 20, 1, 2, 3, 4, 13, 14, 15, 16}}, + {180, {21, 22, 23, 24, 17, 18, 19, 20, 13, 14, 15, 16, + 9, 10, 11, 12, 5, 6, 7, 8, 1, 2, 3, 4}}, + {270, {13, 14, 15, 16, 1, 2, 3, 4, 17, 18, 19, 20, + 5, 6, 7, 8, 21, 22, 23, 24, 9, 10, 11, 12}}}; + + for (auto angle : kAnglesToTest) { + auto output = kOutputBuffers.at(angle).get(); + MP_ASSERT_OK(Rotate(*input, angle, output)); + auto results = kRotationResults.at(angle); + for (int i = 0; i < results.size(); i++) { + EXPECT_EQ(output->plane(0).buffer()[i], results[i]); + } + } +} + +TEST(FrameBufferUtil, RgbRotate) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 3, .height = 2}, + kRotatedDimension = {.width = 2, + .height = 3}; + uint8_t kRgbTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + uint8_t output_data[sizeof(kRgbTestData)]; + auto input = CreateFromRgbRawBuffer(kRgbTestData, kBufferDimension); + const std::array kAnglesToTest = {90, 180, 270}; + std::map> kOutputBuffers; + kOutputBuffers[90] = CreateFromRgbRawBuffer(output_data, kRotatedDimension); + kOutputBuffers[180] = CreateFromRgbRawBuffer(output_data, kBufferDimension); + kOutputBuffers[270] = CreateFromRgbRawBuffer(output_data, kRotatedDimension); + const std::map> kRotationResults{ + {90, {7, 8, 9, 16, 17, 18, 4, 5, 6, 13, 14, 15, 1, 2, 3, 10, 11, 12}}, + {180, {16, 17, 18, 13, 14, 15, 10, 11, 12, 7, 8, 9, 4, 5, 6, 1, 2, 3}}, + {270, {10, 11, 12, 1, 2, 3, 13, 14, 15, 4, 5, 6, 16, 17, 18, 7, 8, 9}}}; + + for (auto angle : kAnglesToTest) { + auto output = kOutputBuffers.at(angle).get(); + MP_ASSERT_OK(Rotate(*input, angle, output)); + auto results = kRotationResults.at(angle); + for (int i = 0; i < results.size(); i++) { + EXPECT_EQ(output->plane(0).buffer()[i], results[i]); + } + } +} + +// Nv21 unit tests. +//------------------------------------------------------------------------------ + +// Helper function to create YUV buffer. +absl::StatusOr> CreateYuvBuffer( + uint8_t* buffer, FrameBuffer::Dimension dimension, int plane_count, + FrameBuffer::Format format) { + DCHECK(plane_count > 0 && plane_count < 4); + ASSIGN_OR_RETURN(auto uv_dimension, GetUvPlaneDimension(dimension, format)); + + if (plane_count == 1) { + const std::vector planes = { + {buffer, /*stride=*/{/*row_stride_bytes=*/dimension.width, + /*pixel_stride_bytes=*/1}}}; + return std::make_shared(planes, dimension, format); + } else if (plane_count == 2) { + CHECK(format == FrameBuffer::Format::kNV12 || + format == FrameBuffer::Format::kNV21); + const std::vector planes = { + {buffer, + /*stride=*/{/*row_stride_bytes=*/dimension.width, + /*pixel_stride_bytes=*/1}}, + {buffer + dimension.Size(), + /*stride=*/{/*row_stride_bytes=*/uv_dimension.width * 2, + /*pixel_stride_bytes=*/2}}}; + return std::make_shared(planes, dimension, format); + } else if (plane_count == 3) { + std::vector planes = { + {buffer, + /*stride=*/{/*row_stride_bytes=*/dimension.width, + /*pixel_stride_bytes=*/1}}, + {buffer + dimension.Size(), + /*stride=*/{/*row_stride_bytes=*/uv_dimension.width, + /*pixel_stride_bytes=*/1}}, + {buffer + dimension.Size() + uv_dimension.Size(), + /*stride=*/{/*row_stride_bytes=*/uv_dimension.width, + /*pixel_stride_bytes=*/1}}}; + return std::make_shared(planes, dimension, format); + } + + return absl::InvalidArgumentError("The plane_count must between 1 and 3."); +} + +TEST(FrameBufferUtil, NV21CreatePlanarYuvBuffer) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}, + kOutputDimension = {.width = 4, .height = 2}; + uint8_t kYTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + uint8_t kUTestData[] = {13, 15, 17, 0, 0, 0}; + uint8_t kVTestData[] = {14, 16, 18, 0, 0, 0}; + uint8_t kNV21VUTestData[] = {14, 13, 16, 15, 18, 17}; + const std::vector three_input_planes = { + {kYTestData, /*stride=*/{6, 1}}, + {kUTestData, /*stride=*/{3, 1}}, + {kVTestData, /*stride=*/{3, 1}}}; + FrameBuffer three_planar_input(three_input_planes, kBufferDimension, + FrameBuffer::Format::kYV21); + + const std::vector two_input_planes = { + {kYTestData, /*stride=*/{6, 1}}, {kNV21VUTestData, /*stride=*/{6, 2}}}; + FrameBuffer two_planar_input(two_input_planes, kBufferDimension, + FrameBuffer::Format::kNV21); + + uint8_t output_y[8], output_u[2], output_v[2]; + const std::vector output_planes = { + {output_y, /*stride=*/{4, 1}}, + {output_u, /*stride=*/{2, 1}}, + {output_v, /*stride=*/{2, 1}}}; + FrameBuffer output(output_planes, kOutputDimension, + FrameBuffer::Format::kYV12); + + MP_ASSERT_OK(Crop(three_planar_input, 2, 0, 5, 1, &output)); + EXPECT_EQ(output.plane(0).buffer()[0], 3); + EXPECT_EQ(output.plane(0).buffer()[1], 4); + EXPECT_EQ(output.plane(0).buffer()[2], 5); + EXPECT_EQ(output.plane(1).buffer()[0], 16); + EXPECT_EQ(output.plane(2).buffer()[0], 15); + + memset(output_y, 0, sizeof(output_y)); + memset(output_u, 0, sizeof(output_u)); + memset(output_v, 0, sizeof(output_v)); + MP_ASSERT_OK(Crop(two_planar_input, 2, 0, 5, 1, &output)); + EXPECT_EQ(output.plane(0).buffer()[0], 3); + EXPECT_EQ(output.plane(0).buffer()[1], 4); + EXPECT_EQ(output.plane(0).buffer()[2], 5); + EXPECT_EQ(output.plane(1).buffer()[0], 16); + EXPECT_EQ(output.plane(2).buffer()[0], 15); +} + +TEST(FrameBufferUtil, NV21Crop) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}, + kOutputDimension = {.width = 4, .height = 2}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN(auto input, + CreateFromRawBuffer(kNV21TestData, kBufferDimension, + FrameBuffer::Format::kNV21)); + uint8_t output_data[12]; + MP_ASSERT_OK_AND_ASSIGN(auto output, + CreateFromRawBuffer(output_data, kOutputDimension, + FrameBuffer::Format::kNV21)); + + MP_ASSERT_OK(Crop(*input, 2, 0, 5, 1, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 3); + EXPECT_EQ(output->plane(0).buffer()[1], 4); + EXPECT_EQ(output->plane(0).buffer()[2], 5); + EXPECT_EQ(output->plane(0).buffer()[8], 15); + EXPECT_EQ(output->plane(0).buffer()[9], 16); +} + +TEST(FrameBufferUtil, YV21Crop) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}, + kOutputDimension = {.width = 4, .height = 2}; + uint8_t kYV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 15, 17, 14, 16, 18}; + MP_ASSERT_OK_AND_ASSIGN( + auto input, + CreateYuvBuffer(kYV21TestData, kBufferDimension, /*plane_count=*/3, + FrameBuffer::Format::kYV21)); + uint8_t output_data[12]{}; + MP_ASSERT_OK_AND_ASSIGN( + auto output, + CreateYuvBuffer(output_data, kOutputDimension, /*plane_count=*/3, + FrameBuffer::Format::kYV21)); + + MP_ASSERT_OK( + Crop(*input, /*x0=*/2, /*y0=*/0, /*x1=*/5, /*y1=*/1, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 3); + EXPECT_EQ(output->plane(0).buffer()[1], 4); + EXPECT_EQ(output->plane(0).buffer()[2], 5); + EXPECT_EQ(output->plane(1).buffer()[0], 15); + EXPECT_EQ(output->plane(2).buffer()[0], 16); +} + +TEST(FrameBufferUtil, NV21HorizontalFlip) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN(auto input, + CreateFromRawBuffer(kNV21TestData, kBufferDimension, + FrameBuffer::Format::kNV21)); + uint8_t output_data[18]; + MP_ASSERT_OK_AND_ASSIGN(auto output, + CreateFromRawBuffer(output_data, kBufferDimension, + FrameBuffer::Format::kNV21)); + + MP_ASSERT_OK(FlipHorizontally(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 6); + EXPECT_EQ(output->plane(0).buffer()[1], 5); + EXPECT_EQ(output->plane(0).buffer()[2], 4); + EXPECT_EQ(output->plane(0).buffer()[12], 17); + EXPECT_EQ(output->plane(0).buffer()[13], 18); +} + +TEST(FrameBufferUtil, NV21VerticalFlip) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN(auto input, + CreateFromRawBuffer(kNV21TestData, kBufferDimension, + FrameBuffer::Format::kNV21)); + uint8_t output_data[18]; + MP_ASSERT_OK_AND_ASSIGN(auto output, + CreateFromRawBuffer(output_data, kBufferDimension, + FrameBuffer::Format::kNV21)); + + MP_ASSERT_OK(FlipVertically(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 7); + EXPECT_EQ(output->plane(0).buffer()[1], 8); + EXPECT_EQ(output->plane(0).buffer()[2], 9); + EXPECT_EQ(output->plane(0).buffer()[12], 13); + EXPECT_EQ(output->plane(0).buffer()[13], 14); +} + +TEST(FrameBufferUtil, NV21Rotate) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}, + kRotatedDimension = {.width = 2, + .height = 6}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN(auto input, + CreateFromRawBuffer(kNV21TestData, kBufferDimension, + FrameBuffer::Format::kNV21)); + uint8_t output_data[18]; + MP_ASSERT_OK_AND_ASSIGN(auto output, + CreateFromRawBuffer(output_data, kRotatedDimension, + FrameBuffer::Format::kNV21)); + + MP_ASSERT_OK(Rotate(*input, 90, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 6); + EXPECT_EQ(output->plane(0).buffer()[1], 12); + EXPECT_EQ(output->plane(0).buffer()[2], 5); + EXPECT_EQ(output->plane(0).buffer()[12], 17); + EXPECT_EQ(output->plane(0).buffer()[13], 18); +} + +TEST(FrameBufferUtil, NV21Resize) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}, + kOutputDimension = {.width = 1, .height = 1}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN(auto input, + CreateFromRawBuffer(kNV21TestData, kBufferDimension, + FrameBuffer::Format::kNV21)); + uint8_t output_data[6]; + MP_ASSERT_OK_AND_ASSIGN(auto output, + CreateFromRawBuffer(output_data, kOutputDimension, + FrameBuffer::Format::kNV21)); + + MP_ASSERT_OK(Resize(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 1); + EXPECT_EQ(output->plane(0).buffer()[1], 13); +} + +TEST(FrameBufferUtil, NV21ConvertGray) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN(auto input, + CreateFromRawBuffer(kNV21TestData, kBufferDimension, + FrameBuffer::Format::kNV21)); + const int kOutputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kGRAY); + std::vector output_data(kOutputSize); + auto output = CreateFromGrayRawBuffer(output_data.data(), kBufferDimension); + + MP_ASSERT_OK(Convert(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 1); + EXPECT_EQ(output->plane(0).buffer()[1], 2); + EXPECT_EQ(output->plane(0).buffer()[11], 12); +} + +TEST(FrameBufferUtil, PaddedYuvConvertGray) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}; + uint8_t kNV21PaddedTestData[] = {1, 2, 3, 4, 5, 6, 100, 100, + 7, 8, 9, 10, 11, 12, 100, 100, + 13, 14, 15, 16, 17, 18, 100, 100}; + constexpr int row_stride_y = 8; + const std::vector planes = { + {kNV21PaddedTestData, /*stride=*/{row_stride_y, 1}}, + {kNV21PaddedTestData + (row_stride_y * kBufferDimension.width), + /*stride=*/{row_stride_y, 2}}}; + auto input = std::make_shared(planes, kBufferDimension, + FrameBuffer::Format::kNV21); + const int kOutputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kGRAY); + std::vector output_data(kOutputSize); + auto output = CreateFromGrayRawBuffer(output_data.data(), kBufferDimension); + + MP_ASSERT_OK(Convert(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 1); + EXPECT_EQ(output->plane(0).buffer()[1], 2); + EXPECT_EQ(output->plane(0).buffer()[6], 7); + EXPECT_EQ(output->plane(0).buffer()[7], 8); + EXPECT_EQ(output->plane(0).buffer()[11], 12); +} + +TEST(FrameBufferUtil, NV21ConvertRgb) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 32, + .height = 8}; + // Note that RGB conversion expects at images width at least width >= 32 + // because the implementation is vectorized. + const int kInputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kNV21); + std::vector input_data(kInputSize); + input_data.data()[0] = 1; + input_data.data()[1] = 2; + input_data.data()[32] = 7; + input_data.data()[33] = 8; + input_data.data()[256] = 13; + input_data.data()[257] = 14; + MP_ASSERT_OK_AND_ASSIGN( + auto input, CreateFromRawBuffer(input_data.data(), kBufferDimension, + FrameBuffer::Format::kNV21)); + const int kOutputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kRGB); + std::vector output_data(kOutputSize); + auto output = CreateFromRgbRawBuffer(output_data.data(), kBufferDimension); + + MP_ASSERT_OK(Convert(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 0); + EXPECT_EQ(output->plane(0).buffer()[1], 122); +} + +TEST(FrameBufferUtil, NV21ConvertHalfRgb) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 64, + .height = 16}, + kOutputDimension = {.width = 32, + .height = 8}; + // Note that RGB conversion expects at images width at least width >= 32 + // because the implementation is vectorized. + uint8_t data[1576]; + for (int i = 0; i < sizeof(data); i++) { + data[i] = (i + 1); + } + MP_ASSERT_OK_AND_ASSIGN( + auto input, + CreateFromRawBuffer(data, kBufferDimension, FrameBuffer::Format::kNV21)); + uint8_t output_data[768]; + auto output = CreateFromRgbRawBuffer(output_data, kOutputDimension); + + MP_ASSERT_OK(Convert(*input, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 0); + EXPECT_EQ(output->plane(0).buffer()[1], 135); +} + +TEST(FrameBufferUtil, NV12ConvertGray) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}; + uint8_t kYTestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + uint8_t kNV12UVTestData[] = {13, 14, 15, 16, 17, 18}; + const std::vector planes_nv12 = { + {kYTestData, /*stride=*/{kBufferDimension.width, 1}}, + {kNV12UVTestData, /*stride=*/{kBufferDimension.width, 2}}}; + auto buffer_nv12 = std::make_shared( + planes_nv12, kBufferDimension, FrameBuffer::Format::kNV12); + const int kOutputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kGRAY); + std::vector output_data(kOutputSize); + auto output = CreateFromGrayRawBuffer(output_data.data(), kBufferDimension); + + MP_ASSERT_OK(Convert(*buffer_nv12, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], kYTestData[0]); + EXPECT_EQ(output->plane(0).buffer()[1], kYTestData[1]); + EXPECT_EQ(output->plane(0).buffer()[11], kYTestData[11]); +} + +TEST(FrameBufferUtil, NV12ConvertRgb) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 32, + .height = 8}; + MP_ASSERT_OK_AND_ASSIGN( + FrameBuffer::Dimension uv_dimension, + GetUvPlaneDimension(kBufferDimension, FrameBuffer::Format::kNV12)); + // Halide RGB converter expects at images width at least width >= 32 because + // the implementation is vectorized. + auto y_data = std::make_unique(kBufferDimension.Size()); + auto uv_data = std::make_unique(uv_dimension.Size() * 2); + y_data[0] = 1; + y_data[1] = 2; + y_data[32] = 7; + y_data[33] = 8; + uv_data[0] = 13; + uv_data[1] = 14; + const std::vector planes_nv12 = { + {y_data.get(), /*stride=*/{kBufferDimension.width, 1}}, + {uv_data.get(), /*stride=*/{kBufferDimension.width, 2}}}; + auto buffer_nv12 = std::make_shared( + planes_nv12, kBufferDimension, FrameBuffer::Format::kNV12); + const int kOutputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kRGB); + std::vector output_data(kOutputSize); + auto output = CreateFromRgbRawBuffer(output_data.data(), kBufferDimension); + + MP_ASSERT_OK(Convert(*buffer_nv12, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 0); + EXPECT_EQ(output->plane(0).buffer()[1], 122); +} + +TEST(FrameBufferUtil, NV12ConvertHalfRgb) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 64, + .height = 16}; + MP_ASSERT_OK_AND_ASSIGN( + FrameBuffer::Dimension uv_dimension, + GetUvPlaneDimension(kBufferDimension, FrameBuffer::Format::kNV12)); + // Halide RGB converter expects at images width at least width >= 32 because + // the implementation is vectorized. + auto y_data = std::make_unique(kBufferDimension.Size()); + auto uv_data = std::make_unique(uv_dimension.Size() * 2); + for (int i = 0; i < kBufferDimension.Size(); i++) { + y_data[i] = (i + 1) % 256; + } + for (int i = 0; i < uv_dimension.Size() * 2; i++) { + uv_data[i] = (i + 1) % 256; + } + const std::vector planes_nv12 = { + {y_data.get(), /*stride=*/{kBufferDimension.width, 1}}, + {uv_data.get(), /*stride=*/{kBufferDimension.width, 2}}}; + auto buffer_nv12 = std::make_shared( + planes_nv12, kBufferDimension, FrameBuffer::Format::kNV12); + constexpr FrameBuffer::Dimension kOutputDimension = { + .width = kBufferDimension.width / 2, + .height = kBufferDimension.height / 2}; + const int kOutputSize = + GetFrameBufferByteSize(kOutputDimension, FrameBuffer::Format::kRGB); + std::vector output_data(kOutputSize); + auto output = CreateFromRgbRawBuffer(output_data.data(), kOutputDimension); + + MP_ASSERT_OK(Convert(*buffer_nv12, output.get())); + EXPECT_EQ(output->plane(0).buffer()[0], 0); + EXPECT_EQ(output->plane(0).buffer()[1], 135); +} + +TEST(FrameBufferUtil, NV21ConvertYV12) { + constexpr FrameBuffer::Dimension kBufferDimension = {.width = 6, .height = 2}; + uint8_t kNV21TestData[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18}; + MP_ASSERT_OK_AND_ASSIGN( + auto nv21, + CreateYuvBuffer(kNV21TestData, kBufferDimension, /*plane_count=*/2, + FrameBuffer::Format::kNV21)); + MP_ASSERT_OK_AND_ASSIGN(FrameBuffer::YuvData nv21_data, + FrameBuffer::GetYuvDataFromFrameBuffer(*nv21)); + const int kOutputSize = + GetFrameBufferByteSize(kBufferDimension, FrameBuffer::Format::kYV12); + std::vector output_data(kOutputSize); + MP_ASSERT_OK_AND_ASSIGN( + auto yv12, + CreateYuvBuffer(output_data.data(), kBufferDimension, /*plane_count=*/3, + FrameBuffer::Format::kYV12)); + MP_ASSERT_OK_AND_ASSIGN(FrameBuffer::YuvData yv12_data, + FrameBuffer::GetYuvDataFromFrameBuffer(*yv12)); + + MP_ASSERT_OK(Convert(*nv21, yv12.get())); + EXPECT_EQ(nv21_data.y_buffer[0], yv12_data.y_buffer[0]); + EXPECT_EQ(nv21_data.u_buffer[0], yv12_data.u_buffer[0]); + EXPECT_EQ(nv21_data.v_buffer[0], yv12_data.v_buffer[0]); +} + +} // namespace +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/gray_buffer.cc b/mediapipe/util/frame_buffer/gray_buffer.cc new file mode 100644 index 000000000..51f7b09e2 --- /dev/null +++ b/mediapipe/util/frame_buffer/gray_buffer.cc @@ -0,0 +1,95 @@ +// 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/gray_buffer.h" + +#include + +#include "mediapipe/util/frame_buffer/buffer_common.h" +#include "mediapipe/util/frame_buffer/halide/gray_flip_halide.h" +#include "mediapipe/util/frame_buffer/halide/gray_resize_halide.h" +#include "mediapipe/util/frame_buffer/halide/gray_rotate_halide.h" +#include "mediapipe/util/frame_buffer/yuv_buffer.h" + +namespace mediapipe { +namespace frame_buffer { + +GrayBuffer::GrayBuffer(uint8_t* buffer, int width, int height) + : owned_buffer_(nullptr) { + Initialize(buffer, width, height); +} + +GrayBuffer::GrayBuffer(int width, int height) + : owned_buffer_(new uint8_t[ByteSize(width, height)]) { + Initialize(owned_buffer_.get(), width, height); +} + +GrayBuffer::GrayBuffer(const GrayBuffer& other) : buffer_(other.buffer_) {} + +GrayBuffer::GrayBuffer(GrayBuffer&& other) { *this = std::move(other); } + +GrayBuffer& GrayBuffer::operator=(const GrayBuffer& other) { + if (this != &other) { + buffer_ = other.buffer_; + } + return *this; +} + +GrayBuffer& GrayBuffer::operator=(GrayBuffer&& other) { + if (this != &other) { + owned_buffer_ = std::move(other.owned_buffer_); + buffer_ = other.buffer_; + } + return *this; +} + +GrayBuffer::~GrayBuffer() {} + +void GrayBuffer::Initialize(uint8_t* data, int width, int height) { + buffer_ = Halide::Runtime::Buffer(data, width, height); +} + +bool GrayBuffer::Crop(int x0, int y0, int x1, int y1) { + // Twiddle the buffer start and extents to crop images. + return common::crop_buffer(x0, y0, x1, y1, buffer()); +} + +bool GrayBuffer::Resize(GrayBuffer* output) { + const int result = gray_resize_halide( + buffer(), static_cast(width()) / output->width(), + static_cast(height()) / output->height(), output->buffer()); + return result == 0; +} + +bool GrayBuffer::Rotate(int angle, GrayBuffer* output) { + const int result = gray_rotate_halide(buffer(), angle, output->buffer()); + return result == 0; +} + +bool GrayBuffer::FlipHorizontally(GrayBuffer* output) { + const int result = gray_flip_halide(buffer(), + false, // horizontal + output->buffer()); + return result == 0; +} + +bool GrayBuffer::FlipVertically(GrayBuffer* output) { + const int result = gray_flip_halide(buffer(), + true, // vertical + output->buffer()); + return result == 0; +} + +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/gray_buffer.h b/mediapipe/util/frame_buffer/gray_buffer.h new file mode 100644 index 000000000..fa181acec --- /dev/null +++ b/mediapipe/util/frame_buffer/gray_buffer.h @@ -0,0 +1,137 @@ +// 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_GRAY_BUFFER_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_GRAY_BUFFER_H_ + +#include + +#include "HalideBuffer.h" +#include "HalideRuntime.h" + +namespace mediapipe { +namespace frame_buffer { + +// GrayBuffer represents a view over a grayscale (i.e. luminance, +// or Y-only) buffer. +// GrayBuffer may be copied and moved efficiently; their backing buffers are +// shared and never deep copied. +// GrayBuffer requires a minimum image width depending on the natural vector +// size of the platform, e.g., 16px. This is not validated by GrayBuffer. +class GrayBuffer { + public: + // Returns the size (in bytes) of a grayscale image of the given + // dimensions. The given dimensions contain padding. + static int ByteSize(int buffer_width, int buffer_height) { + const int size = buffer_width * buffer_height; + return size; + } + + // Builds a grayscale buffer with size as width * height. The buffer should + // be in row-major order with no padding. + // + // Does not take ownership of any backing buffers, which must be large + // enough to fit their contents. + GrayBuffer(uint8_t* buffer, int width, int height); + + // Builds a grayscale buffer with size as width * height. + // + // The underlying backing buffer is allocated and owned by this + // GrayBuffer. + GrayBuffer(int width, int height); + + // GrayBuffer is copyable. The source retains ownership of its backing + // buffers. + // + // Since the source retains ownership of its backing buffer, the source needs + // to outlive this instance's lifetime when the backing buffer is owned by + // the source. Otherwise, the passing in backing buffer should outlive this + // instance. + GrayBuffer(const GrayBuffer& other); + // GrayBuffer is moveable. The source loses ownership of any backing buffers. + // Specifically, if the source owns its backing buffer, after the move, + // Release() will return nullptr. + GrayBuffer(GrayBuffer&& other); + + // GrayBuffer is assignable. The source retains ownership of its backing + // buffers. + // + // Since the source retains ownership of its backing buffer, the source needs + // to outlive this instance's lifetime when the backing buffer is owned by the + // source. Otherwise, the passing in backing buffer should outlive this + // instance. + GrayBuffer& operator=(const GrayBuffer& other); + GrayBuffer& operator=(GrayBuffer&& other); + + ~GrayBuffer(); + + // Performs an in-place crop. Modifies this buffer so that the new extent + // matches that of the given crop rectangle -- (x0, y0) becomes (0, 0) and + // the new width and height are x1 - x0 + 1 and y1 - y0 + 1, respectively. + bool Crop(int x0, int y0, int x1, int y1); + + // Resizes this image to match the dimensions of the given output GrayBuffer + // and places the result into output's backing buffer. + // + // Note, if the output backing buffer is shared with multiple instances, by + // calling this method, all the instances' backing buffers will change. + bool Resize(GrayBuffer* output); + + // Rotates this image into the given buffer by the given angle (90, 180, 270). + // + // Rotation is specified in degrees counter-clockwise such that when rotating + // by 90 degrees, the top-right corner of the source becomes the top-left of + // the output. The output buffer must have its height and width swapped when + // rotating by 90 or 270. + // + // Any angle values other than (90, 180, 270) are invalid. + // + // Note, if the output backing buffer is shared with multiple instances, by + // calling this method, all the instances' backing buffers will change. + bool Rotate(int angle, GrayBuffer* output); + + // Flips this image horizontally/vertically into the given buffer. Both buffer + // dimensions must match. + // + // Note, if the output backing buffer is shared with multiple instances, by + // calling this method, all the instances' backing buffers will change. + bool FlipHorizontally(GrayBuffer* output); + bool FlipVertically(GrayBuffer* output); + + // Releases ownership of the owned backing buffer. + uint8_t* Release() { return owned_buffer_.release(); } + + // Returns the halide_buffer_t* for the image. + halide_buffer_t* buffer() { return buffer_.raw_buffer(); } + + // Returns the image width. + const int width() const { return buffer_.dim(0).extent(); } + // Returns the image height. + const int height() const { return buffer_.dim(1).extent(); } + + private: + void Initialize(uint8_t* data, int width, int height); + + // Non-NULL iff this GrayBuffer owns its buffer. + std::unique_ptr owned_buffer_; + + // Backing buffer: layout is always width x height. The backing buffer binds + // to either "owned_buffer_" or an external buffer. + Halide::Runtime::Buffer buffer_; +}; + +} // namespace frame_buffer +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_GRAY_BUFFER_H_ diff --git a/mediapipe/util/frame_buffer/gray_buffer_test.cc b/mediapipe/util/frame_buffer/gray_buffer_test.cc new file mode 100644 index 000000000..f6f9e9e34 --- /dev/null +++ b/mediapipe/util/frame_buffer/gray_buffer_test.cc @@ -0,0 +1,223 @@ +// 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/gray_buffer.h" + +#include + +#include "absl/log/log.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" + +// The default implementation of halide_error calls abort(), which we don't +// want. Instead, log the error and let the filter invocation fail. +extern "C" void halide_error(void*, const char* message) { + LOG(ERROR) << "Halide Error: " << message; +} + +namespace mediapipe { +namespace frame_buffer { +namespace { + +// Fill a GrayBuffer with zeroes. +void Fill(GrayBuffer* gray_buffer) { + halide_buffer_t* buffer = gray_buffer->buffer(); + for (int y = 0; y < buffer->dim[1].extent; ++y) { + for (int x = 0; x < buffer->dim[0].extent; ++x) { + buffer->host[buffer->dim[1].stride * y + buffer->dim[0].stride * x] = 0; + } + } +} + +TEST(GrayBufferTest, Properties) { + GrayBuffer buffer(5, 4); + EXPECT_EQ(5, buffer.width()); + EXPECT_EQ(4, buffer.height()); +} + +TEST(GrayBufferTest, Release) { + GrayBuffer buffer(4, 4); + delete[] buffer.Release(); +} + +TEST(GrayBufferTest, Assign) { + GrayBuffer buffer(3, 2); + GrayBuffer sink(nullptr, 0, 0); + Fill(&buffer); + sink = buffer; + EXPECT_EQ(3, sink.width()); + EXPECT_EQ(2, sink.height()); + + sink = GrayBuffer(5, 4); + EXPECT_EQ(5, sink.width()); + EXPECT_EQ(4, sink.height()); +} + +TEST(GrayBufferTest, MoveAssign) { + GrayBuffer buffer(3, 2); + GrayBuffer sink(nullptr, 0, 0); + Fill(&buffer); + sink = std::move(buffer); + EXPECT_EQ(nullptr, buffer.Release()); + EXPECT_EQ(3, sink.width()); + EXPECT_EQ(2, sink.height()); +} + +TEST(GrayBufferTest, MoveConstructor) { + GrayBuffer buffer(5, 4); + GrayBuffer sink(std::move(buffer)); + Fill(&buffer); + EXPECT_EQ(nullptr, buffer.Release()); + EXPECT_EQ(5, sink.width()); + EXPECT_EQ(4, sink.height()); +} + +TEST(GrayBufferTest, Crop) { + GrayBuffer source(8, 8); + EXPECT_TRUE(source.Crop(2, 2, 6, 6)); +} + +TEST(GrayBufferTest, Resize_Even) { + uint8_t* data = new uint8_t[16]; + for (int y = 0; y < 4; ++y) { + for (int x = 0; x < 4; ++x) { + data[x + y * 4] = x + y * 4; + } + } + GrayBuffer source(data, 4, 4); + GrayBuffer result(2, 2); + EXPECT_TRUE(source.Resize(&result)); + EXPECT_EQ(0, result.buffer()->host[0]); + EXPECT_EQ(2, result.buffer()->host[1]); + EXPECT_EQ(8, result.buffer()->host[2]); + EXPECT_EQ(10, result.buffer()->host[3]); + delete[] data; +} + +TEST(GrayBufferTest, Resize_Odd) { + uint8_t* data = new uint8_t[16]; + for (int y = 0; y < 4; ++y) { + for (int x = 0; x < 4; ++x) { + data[x + y * 4] = x + y * 4; + } + } + GrayBuffer source(data, 4, 4); + GrayBuffer result(1, 3); + EXPECT_TRUE(source.Resize(&result)); + EXPECT_EQ(0, result.buffer()->host[0]); + EXPECT_EQ(5, result.buffer()->host[1]); + EXPECT_EQ(11, result.buffer()->host[2]); + delete[] data; +} + +TEST(GrayBufferTest, Rotate) { + GrayBuffer buffer(5, 4); + GrayBuffer result(4, 5); + Fill(&buffer); + EXPECT_TRUE(buffer.Rotate(90, &result)); +} + +TEST(GrayBufferTest, Rotate_90) { + uint8_t* data = new uint8_t[4]; + data[0] = 1; + data[1] = 2; + data[2] = 3; + data[3] = 4; + GrayBuffer buffer(data, 2, 2); + GrayBuffer result(2, 2); + EXPECT_TRUE(buffer.Rotate(90, &result)); + + EXPECT_EQ(2, result.buffer()->host[0]); + EXPECT_EQ(4, result.buffer()->host[1]); + EXPECT_EQ(1, result.buffer()->host[2]); + EXPECT_EQ(3, result.buffer()->host[3]); + + delete[] data; +} + +TEST(GrayBufferTest, Rotate_180) { + uint8_t* data = new uint8_t[4]; + data[0] = 1; + data[1] = 2; + data[2] = 3; + data[3] = 4; + GrayBuffer buffer(data, 2, 2); + GrayBuffer result(2, 2); + EXPECT_TRUE(buffer.Rotate(180, &result)); + EXPECT_EQ(4, result.buffer()->host[0]); + EXPECT_EQ(3, result.buffer()->host[1]); + EXPECT_EQ(2, result.buffer()->host[2]); + EXPECT_EQ(1, result.buffer()->host[3]); + delete[] data; +} + +TEST(GrayBufferTest, Rotate_270) { + uint8_t* data = new uint8_t[4]; + data[0] = 1; + data[1] = 2; + data[2] = 3; + data[3] = 4; + GrayBuffer buffer(data, 2, 2); + GrayBuffer result(2, 2); + EXPECT_TRUE(buffer.Rotate(270, &result)); + EXPECT_EQ(3, result.buffer()->host[0]); + EXPECT_EQ(1, result.buffer()->host[1]); + EXPECT_EQ(4, result.buffer()->host[2]); + EXPECT_EQ(2, result.buffer()->host[3]); + delete[] data; +} + +TEST(GrayBufferTest, Flip) { + GrayBuffer buffer(5, 4); + GrayBuffer result(5, 4); + Fill(&buffer); + EXPECT_TRUE(buffer.FlipHorizontally(&result)); + EXPECT_TRUE(buffer.FlipVertically(&result)); +} + +TEST(GrayBufferTest, Flip_Horizontally) { + uint8_t* data = new uint8_t[4]; + data[0] = 1; + data[1] = 2; + data[2] = 3; + data[3] = 4; + GrayBuffer buffer(data, 2, 2); + GrayBuffer result(2, 2); + EXPECT_TRUE(buffer.FlipHorizontally(&result)); + EXPECT_EQ(2, result.buffer()->host[0]); + EXPECT_EQ(1, result.buffer()->host[1]); + EXPECT_EQ(4, result.buffer()->host[2]); + EXPECT_EQ(3, result.buffer()->host[3]); + delete[] data; +} + +TEST(GrayBufferTest, Flip_Vertically) { + uint8_t* data = new uint8_t[4]; + data[0] = 1; + data[1] = 2; + data[2] = 3; + data[3] = 4; + GrayBuffer buffer(data, 2, 2); + GrayBuffer result(2, 2); + EXPECT_TRUE(buffer.FlipVertically(&result)); + EXPECT_EQ(3, result.buffer()->host[0]); + EXPECT_EQ(4, result.buffer()->host[1]); + EXPECT_EQ(1, result.buffer()->host[2]); + EXPECT_EQ(2, result.buffer()->host[3]); + delete[] data; +} + +} // namespace +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/halide/BUILD b/mediapipe/util/frame_buffer/halide/BUILD new file mode 100644 index 000000000..619a16d26 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/BUILD @@ -0,0 +1,118 @@ +# 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. + +load("@halide//:halide.bzl", "halide_library") + +package(default_visibility = ["//mediapipe/util/frame_buffer:__subpackages__"]) + +# Common Halide library: +cc_library( + name = "common", + srcs = ["common.cc"], + hdrs = ["common.h"], + deps = ["@halide//:language"], +) + +# Enable Halide's built-in profiler with: +# bazel ... --define halide_target_features=profile +# and HTML output of its intermediate representation with: +# bazel ... --define halide_extra_outputs=html + +# RGB operations: +halide_library( + name = "rgb_flip_halide", + srcs = ["rgb_flip_generator.cc"], + generator_name = "rgb_flip_generator", +) + +halide_library( + name = "rgb_resize_halide", + srcs = ["rgb_resize_generator.cc"], + generator_deps = [":common"], + generator_name = "rgb_resize_generator", +) + +halide_library( + name = "rgb_rotate_halide", + srcs = ["rgb_rotate_generator.cc"], + generator_deps = [":common"], + generator_name = "rgb_rotate_generator", +) + +halide_library( + name = "rgb_yuv_halide", + srcs = ["rgb_yuv_generator.cc"], + generator_name = "rgb_yuv_generator", +) + +halide_library( + name = "rgb_rgb_halide", + srcs = ["rgb_rgb_generator.cc"], + generator_name = "rgb_rgb_generator", +) + +# YUV operations: +halide_library( + name = "yuv_flip_halide", + srcs = ["yuv_flip_generator.cc"], + generator_name = "yuv_flip_generator", +) + +halide_library( + name = "yuv_rgb_halide", + srcs = ["yuv_rgb_generator.cc"], + generator_name = "yuv_rgb_generator", +) + +halide_library( + name = "yuv_resize_halide", + srcs = ["yuv_resize_generator.cc"], + generator_deps = [":common"], + generator_name = "yuv_resize_generator", +) + +halide_library( + name = "yuv_rotate_halide", + srcs = ["yuv_rotate_generator.cc"], + generator_deps = [":common"], + generator_name = "yuv_rotate_generator", +) + +# Grayscale operations: + +halide_library( + name = "rgb_gray_halide", + srcs = ["rgb_gray_generator.cc"], + generator_name = "rgb_gray_generator", +) + +halide_library( + name = "gray_rotate_halide", + srcs = ["gray_rotate_generator.cc"], + generator_deps = [":common"], + generator_name = "gray_rotate_generator", +) + +halide_library( + name = "gray_flip_halide", + srcs = ["gray_flip_generator.cc"], + generator_name = "gray_flip_generator", +) + +halide_library( + name = "gray_resize_halide", + srcs = ["gray_resize_generator.cc"], + generator_deps = [":common"], + generator_name = "gray_resize_generator", +) diff --git a/mediapipe/util/frame_buffer/halide/common.cc b/mediapipe/util/frame_buffer/halide/common.cc new file mode 100644 index 000000000..8142c82f5 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/common.cc @@ -0,0 +1,89 @@ +// 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/halide/common.h" + +namespace mediapipe { +namespace frame_buffer { +namespace halide { +namespace common { + +namespace { +using ::Halide::_; +} + +void resize_nn(Halide::Func input, Halide::Func result, Halide::Expr fx, + Halide::Expr fy) { + Halide::Var x{"x"}, y{"y"}; + result(x, y, _) = input(Halide::cast((x + 0.5f) * fx), + Halide::cast((y + 0.5f) * fy), _); +} + +// Borrowed from photos/editing/halide/src/resize_image_bilinear_generator.cc: +void resize_bilinear(Halide::Func input, Halide::Func result, Halide::Expr fx, + Halide::Expr fy) { + Halide::Var x{"x"}, y{"y"}; + Halide::Func x_interpolated("x_interpolated"); + + Halide::Expr xi = Halide::cast(x * fx); + Halide::Expr xr = x * fx - xi; + Halide::Expr x0 = input(xi + 0, y, _); + Halide::Expr x1 = input(xi + 1, y, _); + x_interpolated(x, y, _) = lerp(x0, x1, xr); + + Halide::Expr yi = Halide::cast(y * fy); + Halide::Expr yr = y * fy - yi; + Halide::Expr y0 = x_interpolated(x, yi + 0, _); + Halide::Expr y1 = x_interpolated(x, yi + 1, _); + result(x, y, _) = lerp(y0, y1, yr); +} + +void resize_bilinear_int(Halide::Func input, Halide::Func result, + Halide::Expr fx, Halide::Expr fy) { + Halide::Var x{"x"}, y{"y"}; + Halide::Func x_interpolated("x_interpolated"); + + fx = Halide::cast(fx * 65536); + Halide::Expr xi = Halide::cast(x * fx / 65536); + Halide::Expr xr = Halide::cast(x * fx % 65536); + Halide::Expr x0 = input(xi + 0, y, _); + Halide::Expr x1 = input(xi + 1, y, _); + x_interpolated(x, y, _) = lerp(x0, x1, xr); + + fy = Halide::cast(fy * 65536); + Halide::Expr yi = Halide::cast(y * fy / 65536); + Halide::Expr yr = Halide::cast(y * fy % 65536); + Halide::Expr y0 = x_interpolated(x, yi + 0, _); + Halide::Expr y1 = x_interpolated(x, yi + 1, _); + result(x, y, _) = lerp(y0, y1, yr); +} + +void rotate(Halide::Func input, Halide::Func result, Halide::Expr width, + Halide::Expr height, Halide::Expr angle) { + Halide::Var x{"x"}, y{"y"}; + Halide::Func result_90_degrees, result_180_degrees, result_270_degrees; + result_90_degrees(x, y, _) = input(width - 1 - y, x, _); + result_180_degrees(x, y, _) = input(width - 1 - x, height - 1 - y, _); + result_270_degrees(x, y, _) = input(y, height - 1 - x, _); + + result(x, y, _) = + select(angle == 90, result_90_degrees(x, y, _), angle == 180, + result_180_degrees(x, y, _), angle == 270, + result_270_degrees(x, y, _), input(x, y, _)); +} + +} // namespace common +} // namespace halide +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/halide/common.h b/mediapipe/util/frame_buffer/halide/common.h new file mode 100644 index 000000000..4b7023a9f --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/common.h @@ -0,0 +1,61 @@ +// 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_HALIDE_COMMON_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_HALIDE_COMMON_H_ + +#include "Halide.h" + +namespace mediapipe { +namespace frame_buffer { +namespace halide { +namespace common { + +template +Halide::Expr is_planar(const T& buffer) { + return buffer.dim(0).stride() == 1; +} + +template +Halide::Expr is_interleaved(const T& buffer) { + return buffer.dim(0).stride() == buffer.dim(2).extent() && + buffer.dim(2).stride() == 1; +} + +// Resize scale parameters (fx, fy) are the ratio of source size to output +// size; thus if you want to produce an image half as wide and twice as tall +// as the input, (fx, fy) should be (2, 0.5). + +// Nearest-neighbor resize: fast, but low-quality (prone to aliasing). +void resize_nn(Halide::Func input, Halide::Func result, Halide::Expr fx, + Halide::Expr fy); + +// Resize with bilinear interpolation: slower but higher-quality. +void resize_bilinear(Halide::Func input, Halide::Func result, Halide::Expr fx, + Halide::Expr fy); +// Identical to the above, except that it uses fixed point integer math. +void resize_bilinear_int(Halide::Func input, Halide::Func result, + Halide::Expr fx, Halide::Expr fy); + +// Note: width and height are the source image dimensions; angle must be one +// of [0, 90, 180, 270] or the result is undefined. +void rotate(Halide::Func input, Halide::Func result, Halide::Expr width, + Halide::Expr height, Halide::Expr angle); + +} // namespace common +} // namespace halide +} // namespace frame_buffer +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_HALIDE_COMMON_H_ diff --git a/mediapipe/util/frame_buffer/halide/gray_flip_generator.cc b/mediapipe/util/frame_buffer/halide/gray_flip_generator.cc new file mode 100644 index 000000000..528489c84 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/gray_flip_generator.cc @@ -0,0 +1,61 @@ +// 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 { + +using ::Halide::_; + +class GrayFlip : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_y{"src_y"}; + + // Flip vertically if true; flips horizontally (mirroring) otherwise. + Input flip_vertical{"flip_vertical", false}; + + Output dst_y{"dst_y", UInt(8), 2}; + + void generate(); + void schedule(); + + private: + void flip(Func input, Func result, Expr width, Expr height, Expr vertical); +}; + +void GrayFlip::generate() { + Halide::Func flip_x, flip_y; + flip_x(x, y, _) = src_y(src_y.dim(0).extent() - x - 1, y, _); + flip_y(x, y, _) = src_y(x, src_y.dim(1).extent() - y - 1, _); + + dst_y(x, y, _) = select(flip_vertical, flip_y(x, y, _), flip_x(x, y, _)); +} + +void GrayFlip::schedule() { + Halide::Func dst_y_func = dst_y; + + // Y plane dimensions start at zero and destination bounds must match. + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + dst_y_output.dim(0).set_bounds(0, src_y.dim(0).extent()); + dst_y_output.dim(1).set_bounds(0, src_y.dim(1).extent()); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(GrayFlip, gray_flip_generator) diff --git a/mediapipe/util/frame_buffer/halide/gray_resize_generator.cc b/mediapipe/util/frame_buffer/halide/gray_resize_generator.cc new file mode 100644 index 000000000..adda9c8d5 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/gray_resize_generator.cc @@ -0,0 +1,60 @@ +// 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" +#include "mediapipe/util/frame_buffer/halide/common.h" + +namespace { + +using ::Halide::BoundaryConditions::repeat_edge; +using ::mediapipe::frame_buffer::halide::common::resize_bilinear_int; + +class GrayResize : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + Input> src_y{"src_y"}; + Input scale_x{"scale_x", 1.0f, 0.0f, 1024.0f}; + Input scale_y{"scale_y", 1.0f, 0.0f, 1024.0f}; + + Output dst_y{"dst_y", UInt(8), 2}; + + void generate(); + void schedule(); +}; + +void GrayResize::generate() { + resize_bilinear_int(repeat_edge(src_y), dst_y, scale_x, scale_y); +} + +void GrayResize::schedule() { + // Grayscale image dimensions start at zero. + Halide::Func dst_y_func = dst_y; + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + dst_y_output.dim(0).set_min(0); + dst_y_output.dim(1).set_min(0); + + // We must ensure that the image is wide enough to support vector + // operations. + const int vector_size = natural_vector_size(); + Halide::Expr min_y_width = + Halide::min(src_y.dim(0).extent(), dst_y_output.dim(0).extent()); + dst_y_func.specialize(min_y_width >= vector_size).vectorize(x, vector_size); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(GrayResize, gray_resize_generator) diff --git a/mediapipe/util/frame_buffer/halide/gray_rotate_generator.cc b/mediapipe/util/frame_buffer/halide/gray_rotate_generator.cc new file mode 100644 index 000000000..825741a5f --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/gray_rotate_generator.cc @@ -0,0 +1,63 @@ +// 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" +#include "mediapipe/util/frame_buffer/halide/common.h" + +namespace { + +using ::mediapipe::frame_buffer::halide::common::rotate; + +class GrayRotate : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_y{"src_y"}; + + // Rotation angle in degrees counter-clockwise. Must be in {0, 90, 180, 270}. + Input rotation_angle{"rotation_angle", 0}; + + Output dst_y{"dst_y", UInt(8), 2}; + + void generate(); + void schedule(); +}; + +void GrayRotate::generate() { + const Halide::Expr width = src_y.dim(0).extent(); + const Halide::Expr height = src_y.dim(1).extent(); + + rotate(src_y, dst_y, width, height, rotation_angle); +} + +void GrayRotate::schedule() { + Halide::Func dst_y_func = dst_y; + dst_y_func.specialize(rotation_angle == 0).reorder(x, y); + dst_y_func.specialize(rotation_angle == 90).reorder(y, x); + dst_y_func.specialize(rotation_angle == 180).reorder(x, y); + dst_y_func.specialize(rotation_angle == 270).reorder(y, x); + + // Y plane dimensions start at zero. We could additionally constrain the + // extent to be even, but that doesn't seem to have any benefit. + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + dst_y_output.dim(0).set_min(0); + dst_y_output.dim(1).set_min(0); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(GrayRotate, gray_rotate_generator) diff --git a/mediapipe/util/frame_buffer/halide/rgb_flip_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_flip_generator.cc new file mode 100644 index 000000000..2b2723a68 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_flip_generator.cc @@ -0,0 +1,84 @@ +// 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 { + +using ::Halide::_; + +class RgbFlip : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_rgb{"src_rgb"}; + // Flip vertically if true; flips horizontally (mirroring) otherwise. + Input flip_vertical{"flip_vertical", false}; + + Output dst_rgb{"dst_rgb", UInt(8), 3}; + + void generate(); + void schedule(); + + private: + void flip(Func input, Func result, Expr width, Expr height, Expr vertical); +}; + +void RgbFlip::flip(Halide::Func input, Halide::Func result, Halide::Expr width, + Halide::Expr height, Halide::Expr vertical) { + Halide::Func flip_x, flip_y; + flip_x(x, y, _) = input(width - x - 1, y, _); + flip_y(x, y, _) = input(x, height - y - 1, _); + + result(x, y, _) = select(vertical, flip_y(x, y, _), flip_x(x, y, _)); +} + +void RgbFlip::generate() { + const Halide::Expr width = src_rgb.dim(0).extent(); + const Halide::Expr height = src_rgb.dim(1).extent(); + + // Flip each of the RGB planes independently. + flip(src_rgb, dst_rgb, width, height, flip_vertical); +} + +void RgbFlip::schedule() { + Halide::Func dst_rgb_func = dst_rgb; + Halide::Var c = dst_rgb_func.args()[2]; + Halide::OutputImageParam rgb_output = dst_rgb_func.output_buffer(); + + // Iterate over channel in the innermost loop, then x, then y. + dst_rgb_func.reorder(c, x, y); + + // RGB planes starts at index zero in every dimension and destination bounds + // must match. + src_rgb.dim(0).set_min(0); + src_rgb.dim(1).set_min(0); + src_rgb.dim(2).set_min(0); + rgb_output.dim(0).set_bounds(0, src_rgb.dim(0).extent()); + rgb_output.dim(1).set_bounds(0, src_rgb.dim(1).extent()); + rgb_output.dim(2).set_bounds(0, src_rgb.dim(2).extent()); + + // Require that the input/output buffer be interleaved and tightly- + // packed; that is, either RGBRGBRGB[...] or RGBARGBARGBA[...], + // without gaps between pixels. + src_rgb.dim(0).set_stride(src_rgb.dim(2).extent()); + src_rgb.dim(2).set_stride(1); + rgb_output.dim(0).set_stride(rgb_output.dim(2).extent()); + rgb_output.dim(2).set_stride(1); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbFlip, rgb_flip_generator) diff --git a/mediapipe/util/frame_buffer/halide/rgb_gray_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_gray_generator.cc new file mode 100644 index 000000000..1df96540a --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_gray_generator.cc @@ -0,0 +1,65 @@ +// 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 RgbGray : public Halide::Generator { + public: + Var x{"x"}, y{"y"}, c{"c"}; + + Input> src_rgb{"rgb"}; + Output> convert{"convert"}; + + void generate(); + void schedule(); +}; + +// Integer math versions of the full-range JFIF RGB-Y coefficients. +// Y = 0.2990*R + 0.5870*G + 0.1140*B +// See https://www.w3.org/Graphics/JPEG/jfif3.pdf. These coefficients are +// similar to, but not identical, to those used in Android. +Halide::Expr rgby(Halide::Expr r, Halide::Expr g, Halide::Expr b) { + r = Halide::cast(r); + g = Halide::cast(g); + b = Halide::cast(b); + return (19595 * r + 38470 * g + 7474 * b + 32768) >> 16; +} + +void RgbGray::generate() { + Halide::Func gray("gray"); + gray(x, y) = rgby(src_rgb(x, y, 0), src_rgb(x, y, 1), src_rgb(x, y, 2)); + convert(x, y) = Halide::saturating_cast(gray(x, y)); +} + +void RgbGray::schedule() { + // RGB images starts at index zero in every dimension. + src_rgb.dim(0).set_min(0); + src_rgb.dim(1).set_min(0); + src_rgb.dim(2).set_min(0); + + // Require that the input buffer be interleaved and tightly-packed; + // with no gaps between pixels. + src_rgb.dim(0).set_stride(src_rgb.dim(2).extent()); + src_rgb.dim(2).set_stride(1); + + // Grayscale images starts at index zero in every dimension. + convert.dim(0).set_min(0); + convert.dim(1).set_min(0); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbGray, rgb_gray_generator) diff --git a/mediapipe/util/frame_buffer/halide/rgb_resize_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_resize_generator.cc new file mode 100644 index 000000000..469120ec3 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_resize_generator.cc @@ -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. + +#include "Halide.h" +#include "mediapipe/util/frame_buffer/halide/common.h" + +namespace { + +using ::Halide::BoundaryConditions::repeat_edge; +using ::mediapipe::frame_buffer::halide::common::resize_bilinear_int; + +class RgbResize : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + Input> src_rgb{"src_rgb"}; + Input scale_x{"scale_x", 1.0f, 0.0f, 1024.0f}; + Input scale_y{"scale_y", 1.0f, 0.0f, 1024.0f}; + + Output dst_rgb{"dst_rgb", UInt(8), 3}; + + void generate(); + void schedule(); +}; + +void RgbResize::generate() { + // Resize each of the RGB planes independently. + resize_bilinear_int(repeat_edge(src_rgb), dst_rgb, scale_x, scale_y); +} + +void RgbResize::schedule() { + Halide::Func dst_rgb_func = dst_rgb; + Halide::Var c = dst_rgb_func.args()[2]; + Halide::OutputImageParam rgb_output = dst_rgb_func.output_buffer(); + Halide::Expr input_rgb_channels = src_rgb.dim(2).extent(); + Halide::Expr output_rgb_channels = rgb_output.dim(2).extent(); + Halide::Expr min_width = + Halide::min(src_rgb.dim(0).extent(), rgb_output.dim(0).extent()); + + // Specialize the generated code for RGB and RGBA (input and output channels + // must match); further, specialize the vectorized implementation so it only + // runs on images wide enough to support it. + const int vector_size = natural_vector_size(); + const Expr channel_specializations[] = { + input_rgb_channels == 3 && output_rgb_channels == 3, + input_rgb_channels == 4 && output_rgb_channels == 4, + }; + dst_rgb_func.reorder(c, x, y); + for (const Expr& channel_specialization : channel_specializations) { + dst_rgb_func.specialize(channel_specialization && min_width >= vector_size) + .unroll(c) + .vectorize(x, vector_size); + } + + // Require that the input/output buffer be interleaved and tightly- + // packed; that is, either RGBRGBRGB[...] or RGBARGBARGBA[...], + // without gaps between pixels. + src_rgb.dim(0).set_stride(input_rgb_channels); + src_rgb.dim(2).set_stride(1); + rgb_output.dim(0).set_stride(output_rgb_channels); + rgb_output.dim(2).set_stride(1); + + // RGB planes starts at index zero in every dimension. + src_rgb.dim(0).set_min(0); + src_rgb.dim(1).set_min(0); + src_rgb.dim(2).set_min(0); + rgb_output.dim(0).set_min(0); + rgb_output.dim(1).set_min(0); + rgb_output.dim(2).set_min(0); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbResize, rgb_resize_generator) diff --git a/mediapipe/util/frame_buffer/halide/rgb_rgb_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_rgb_generator.cc new file mode 100644 index 000000000..99a016896 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_rgb_generator.cc @@ -0,0 +1,64 @@ +// 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 + +#include "Halide.h" + +namespace { + +// Convert rgb_buffer between 3 and 4 channels. When converting from 3 channels +// to 4 channels, the alpha value is always 255. +class RgbRgb : public Halide::Generator { + public: + Var x{"x"}, y{"y"}, c{"c"}; + + Input> src_rgb{"src_rgb"}; + Output> dst_rgb{"dst_rgb"}; + + void generate(); + void schedule(); +}; + +void RgbRgb::generate() { + // We use Halide::clamp to avoid evaluating src_rgb(x, y, c) when c == 3 and + // the src_rgb only has c <= 2 (rgb -> rgba conversion case). + dst_rgb(x, y, c) = + Halide::select(c == 3, 255, src_rgb(x, y, Halide::clamp(c, 0, 2))); +} + +void RgbRgb::schedule() { + Halide::Expr input_rgb_channels = src_rgb.dim(2).extent(); + Halide::Expr output_rgb_channels = dst_rgb.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_rgb.dim(0).set_min(0); + dst_rgb.dim(1).set_min(0); + dst_rgb.dim(2).set_min(0); + dst_rgb.dim(0).set_stride(output_rgb_channels); + dst_rgb.dim(2).set_stride(1); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbRgb, rgb_rgb_generator) diff --git a/mediapipe/util/frame_buffer/halide/rgb_rotate_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_rotate_generator.cc new file mode 100644 index 000000000..aa2bb24ec --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_rotate_generator.cc @@ -0,0 +1,76 @@ +// 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" +#include "mediapipe/util/frame_buffer/halide/common.h" + +namespace { + +using ::mediapipe::frame_buffer::halide::common::rotate; + +class RgbRotate : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_rgb{"src_rgb"}; + // Rotation angle in degrees counter-clockwise. Must be in {0, 90, 180, 270}. + Input rotation_angle{"rotation_angle", 0}; + + Output dst_rgb{"dst_rgb", UInt(8), 3}; + + void generate(); + void schedule(); +}; + +void RgbRotate::generate() { + const Halide::Expr width = src_rgb.dim(0).extent(); + const Halide::Expr height = src_rgb.dim(1).extent(); + + // Rotate each of the RGB planes independently. + rotate(src_rgb, dst_rgb, width, height, rotation_angle); +} + +void RgbRotate::schedule() { + // TODO: Remove specialization for (angle == 0) since that is + // a no-op and callers should simply skip rotation. Doing so would cause + // a bounds assertion crash if called with angle=0, however. + Halide::Func dst_rgb_func = dst_rgb; + Halide::Var c = dst_rgb_func.args()[2]; + Halide::OutputImageParam rgb_output = dst_rgb_func.output_buffer(); + dst_rgb_func.specialize(rotation_angle == 0).reorder(c, x, y); + dst_rgb_func.specialize(rotation_angle == 90).reorder(c, y, x); + dst_rgb_func.specialize(rotation_angle == 180).reorder(c, x, y); + dst_rgb_func.specialize(rotation_angle == 270).reorder(c, y, x); + + // RGB planes starts at index zero in every dimension. + src_rgb.dim(0).set_min(0); + src_rgb.dim(1).set_min(0); + src_rgb.dim(2).set_min(0); + rgb_output.dim(0).set_min(0); + rgb_output.dim(1).set_min(0); + rgb_output.dim(2).set_min(0); + + // Require that the input/output buffer be interleaved and tightly- + // packed; that is, either RGBRGBRGB[...] or RGBARGBARGBA[...], + // without gaps between pixels. + src_rgb.dim(0).set_stride(src_rgb.dim(2).extent()); + src_rgb.dim(2).set_stride(1); + rgb_output.dim(0).set_stride(rgb_output.dim(2).extent()); + rgb_output.dim(2).set_stride(1); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbRotate, rgb_rotate_generator) diff --git a/mediapipe/util/frame_buffer/halide/rgb_yuv_generator.cc b/mediapipe/util/frame_buffer/halide/rgb_yuv_generator.cc new file mode 100644 index 000000000..283dcec8a --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/rgb_yuv_generator.cc @@ -0,0 +1,101 @@ +// 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 RgbYuv : public Halide::Generator { + public: + Var x{"x"}, y{"y"}, c{"c"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_rgb{"rgb"}; + + Output dst_y{"dst_y", UInt(8), 2}; + Output dst_uv{"dst_uv", UInt(8), 3}; + + void generate(); + void schedule(); +}; + +// Integer math versions of the full-range JFIF RGB-YUV coefficients. +// Y = 0.2990*R + 0.5870*G + 0.1140*B +// U = -0.1687*R - 0.3313*G + 0.5000*B + 128 +// V = 0.5000*R - 0.4187*G - 0.0813*B + 128 +// See https://www.w3.org/Graphics/JPEG/jfif3.pdf. These coefficients are +// similar to, but not identical, to those used in Android. +Halide::Tuple rgbyuv(Halide::Expr r, Halide::Expr g, Halide::Expr b) { + r = Halide::cast(r); + g = Halide::cast(g); + b = Halide::cast(b); + return { + (19595 * r + 38470 * g + 7474 * b + 32768) >> 16, + ((-11056 * r - 21712 * g + 32768 * b + 32768) >> 16) + 128, + ((32768 * r - 27440 * g - 5328 * b + 32768) >> 16) + 128, + }; +} + +void RgbYuv::generate() { + Halide::Func yuv_tuple("yuv_tuple"); + yuv_tuple(x, y) = + rgbyuv(src_rgb(x, y, 0), src_rgb(x, y, 1), src_rgb(x, y, 2)); + + // Y values are copied one-for-one; UV values are sampled 1/4. + // TODO: Take the average UV values across the 2x2 block. + dst_y(x, y) = Halide::saturating_cast(yuv_tuple(x, y)[0]); + dst_uv(x, y, c) = Halide::saturating_cast(Halide::select( + c == 0, yuv_tuple(x * 2, y * 2)[2], yuv_tuple(x * 2, y * 2)[1])); + // NOTE: uv channel indices above assume NV21; this can be abstracted out + // by twiddling strides in calling code. +} + +void RgbYuv::schedule() { + // RGB images starts at index zero in every dimension. + src_rgb.dim(0).set_min(0); + src_rgb.dim(1).set_min(0); + src_rgb.dim(2).set_min(0); + + // Require that the input buffer be interleaved and tightly-packed; + // that is, either RGBRGBRGB[...] or RGBARGBARGBA[...], without gaps + // between pixels. + src_rgb.dim(0).set_stride(src_rgb.dim(2).extent()); + src_rgb.dim(2).set_stride(1); + + // Y plane dimensions start at zero. We could additionally constrain the + // extent to be even, but that doesn't seem to have any benefit. + Halide::Func dst_y_func = dst_y; + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + dst_y_output.dim(0).set_min(0); + dst_y_output.dim(1).set_min(0); + + // UV plane has two channels and is half the size of the Y plane in X/Y. + Halide::Func dst_uv_func = dst_uv; + Halide::OutputImageParam dst_uv_output = dst_uv_func.output_buffer(); + dst_uv_output.dim(0).set_bounds(0, (dst_y_output.dim(0).extent() + 1) / 2); + dst_uv_output.dim(1).set_bounds(0, (dst_y_output.dim(1).extent() + 1) / 2); + dst_uv_output.dim(2).set_bounds(0, 2); + + // UV channel processing should be loop unrolled. + dst_uv_func.reorder(c, x, y); + dst_uv_func.unroll(c); + + // Remove default memory layout constraints and accept/produce generic UV + // (including semi-planar and planar). + dst_uv_output.dim(0).set_stride(Expr()); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(RgbYuv, rgb_yuv_generator) diff --git a/mediapipe/util/frame_buffer/halide/yuv_flip_generator.cc b/mediapipe/util/frame_buffer/halide/yuv_flip_generator.cc new file mode 100644 index 000000000..83080f3d7 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/yuv_flip_generator.cc @@ -0,0 +1,90 @@ +// 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 { + +using ::Halide::_; + +class YuvFlip : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_y{"src_y"}; + Input> src_uv{"src_uv"}; + // Flip vertically if true; flips horizontally (mirroring) otherwise. + Input flip_vertical{"flip_vertical", false}; + + Output dst_y{"dst_y", UInt(8), 2}; + Output dst_uv{"dst_uv", UInt(8), 3}; + + void generate(); + void schedule(); + + private: + void flip(Func input, Func result, Expr width, Expr height, Expr vertical); +}; + +void YuvFlip::flip(Halide::Func input, Halide::Func result, Halide::Expr width, + Halide::Expr height, Halide::Expr vertical) { + Halide::Func flip_x, flip_y; + flip_x(x, y, _) = input(width - x - 1, y, _); + flip_y(x, y, _) = input(x, height - y - 1, _); + + result(x, y, _) = select(vertical, flip_y(x, y, _), flip_x(x, y, _)); +} + +void YuvFlip::generate() { + const Halide::Expr width = src_y.dim(0).extent(); + const Halide::Expr height = src_y.dim(1).extent(); + + // Flip each of the YUV planes independently. + flip(src_y, dst_y, width, height, flip_vertical); + flip(src_uv, dst_uv, (width + 1) / 2, (height + 1) / 2, flip_vertical); +} + +void YuvFlip::schedule() { + Halide::Func dst_y_func = dst_y; + Halide::Func dst_uv_func = dst_uv; + Halide::Var c = dst_uv_func.args()[2]; + dst_uv_func.unroll(c); + dst_uv_func.reorder(c, x, y); + + // Y plane dimensions start at zero and destination bounds must match. + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + dst_y_output.dim(0).set_bounds(0, src_y.dim(0).extent()); + dst_y_output.dim(1).set_bounds(0, src_y.dim(1).extent()); + + // UV plane has two channels and is half the size of the Y plane in X/Y. + Halide::OutputImageParam dst_uv_output = dst_uv_func.output_buffer(); + src_uv.dim(0).set_bounds(0, (src_y.dim(0).extent() + 1) / 2); + src_uv.dim(1).set_bounds(0, (src_y.dim(1).extent() + 1) / 2); + src_uv.dim(2).set_bounds(0, 2); + dst_uv_output.dim(0).set_bounds(0, (dst_y_output.dim(0).extent() + 1) / 2); + dst_uv_output.dim(1).set_bounds(0, (dst_y_output.dim(1).extent() + 1) / 2); + dst_uv_output.dim(2).set_bounds(0, 2); + + // Remove default memory layout constraints and accept/produce generic UV + // (including semi-planar and planar). + src_uv.dim(0).set_stride(Expr()); + dst_uv_output.dim(0).set_stride(Expr()); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(YuvFlip, yuv_flip_generator) diff --git a/mediapipe/util/frame_buffer/halide/yuv_resize_generator.cc b/mediapipe/util/frame_buffer/halide/yuv_resize_generator.cc new file mode 100644 index 000000000..805877fca --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/yuv_resize_generator.cc @@ -0,0 +1,91 @@ +// 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" +#include "mediapipe/util/frame_buffer/halide/common.h" + +namespace { + +using ::Halide::BoundaryConditions::repeat_edge; +using ::mediapipe::frame_buffer::halide::common::is_interleaved; +using ::mediapipe::frame_buffer::halide::common::is_planar; +using ::mediapipe::frame_buffer::halide::common::resize_bilinear_int; + +class YuvResize : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + Input> src_y{"src_y"}; + Input> src_uv{"src_uv"}; + Input scale_x{"scale_x", 1.0f, 0.0f, 1024.0f}; + Input scale_y{"scale_y", 1.0f, 0.0f, 1024.0f}; + + Output dst_y{"dst_y", UInt(8), 2}; + Output dst_uv{"dst_uv", UInt(8), 3}; + + void generate(); + void schedule(); +}; + +void YuvResize::generate() { + // Resize each of the YUV planes independently. + resize_bilinear_int(repeat_edge(src_y), dst_y, scale_x, scale_y); + resize_bilinear_int(repeat_edge(src_uv), dst_uv, scale_x, scale_y); +} + +void YuvResize::schedule() { + // Y plane dimensions start at zero. We could additionally constrain the + // extent to be even, but that doesn't seem to have any benefit. + Halide::Func dst_y_func = dst_y; + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + dst_y_output.dim(0).set_min(0); + dst_y_output.dim(1).set_min(0); + + // UV plane has two channels and is half the size of the Y plane in X/Y. + Halide::Func dst_uv_func = dst_uv; + Halide::OutputImageParam dst_uv_output = dst_uv_func.output_buffer(); + src_uv.dim(0).set_bounds(0, (src_y.dim(0).extent() + 1) / 2); + src_uv.dim(1).set_bounds(0, (src_y.dim(1).extent() + 1) / 2); + src_uv.dim(2).set_bounds(0, 2); + dst_uv_output.dim(0).set_bounds(0, (dst_y_output.dim(0).extent() + 1) / 2); + dst_uv_output.dim(1).set_bounds(0, (dst_y_output.dim(1).extent() + 1) / 2); + dst_uv_output.dim(2).set_bounds(0, 2); + + // With bilinear filtering enabled, Y plane resize is profitably vectorizable + // though we must ensure that the image is wide enough to support vector + // operations. + const int vector_size = natural_vector_size(); + Halide::Expr min_y_width = + Halide::min(src_y.dim(0).extent(), dst_y_output.dim(0).extent()); + dst_y_func.specialize(min_y_width >= vector_size).vectorize(x, vector_size); + + // Remove default memory layout constraints and generate specialized + // fast-path implementations when both UV source and output are either + // planar or interleaved. Everything else falls onto a slow path. + src_uv.dim(0).set_stride(Expr()); + dst_uv_output.dim(0).set_stride(Expr()); + + Halide::Var c = dst_uv_func.args()[2]; + dst_uv_func + .specialize(is_interleaved(src_uv) && is_interleaved(dst_uv_output)) + .reorder(c, x, y) + .unroll(c); + dst_uv_func.specialize(is_planar(src_uv) && is_planar(dst_uv_output)); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(YuvResize, yuv_resize_generator) diff --git a/mediapipe/util/frame_buffer/halide/yuv_rgb_generator.cc b/mediapipe/util/frame_buffer/halide/yuv_rgb_generator.cc new file mode 100644 index 000000000..916ea290a --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/yuv_rgb_generator.cc @@ -0,0 +1,110 @@ +// 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 YuvRgb : public Halide::Generator { + public: + Var x{"x"}, y{"y"}, c{"c"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_y{"src_y"}; + Input> src_uv{"src_uv"}; + Input halve{"halve", false}; + + Output rgb{"rgb", UInt(8), 3}; + + void generate(); + void schedule(); +}; + +Halide::Expr demux(Halide::Expr c, Halide::Tuple values) { + return select(c == 0, values[0], c == 1, values[1], c == 2, values[2], 255); +} + +// Integer math versions of the full-range JFIF YUV-RGB coefficients. +// R = Y' + 1.40200*(V-128) +// G = Y' - 0.34414*(U-128) - 0.71414*(V-128) +// B = Y' + 1.77200*(U-128) +// See https://www.w3.org/Graphics/JPEG/jfif3.pdf. These coefficients are +// similar to, but not identical, to those used in Android. +Halide::Tuple yuvrgb(Halide::Expr y, Halide::Expr u, Halide::Expr v) { + y = Halide::cast(y); + u = Halide::cast(u) - 128; + v = Halide::cast(v) - 128; + return { + y + ((91881 * v + 32768) >> 16), + y - ((22544 * u + 46802 * v + 32768) >> 16), + y + ((116130 * u + 32768) >> 16), + }; +} + +void YuvRgb::generate() { + // Each 2x2 block of Y pixels shares the same UV values, so UV-coordinates + // advance half as slowly as Y-coordinates. When taking advantage of the + // "free" 2x downsampling, use every UV value but skip every other Y. + Halide::Expr yx = select(halve, 2 * x, x), yy = select(halve, 2 * y, y); + Halide::Expr uvx = select(halve, x, x / 2), uvy = select(halve, y, y / 2); + + rgb(x, y, c) = Halide::saturating_cast(demux( + c, yuvrgb(src_y(yx, yy), src_uv(uvx, uvy, 1), src_uv(uvx, uvy, 0)))); + // NOTE: uv channel indices above assume NV21; this can be abstracted out + // by twiddling strides in calling code. +} + +void YuvRgb::schedule() { + // Y plane dimensions start at zero. We could additionally constrain the + // extent to be even, but that doesn't seem to have any benefit. + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + + // UV plane has two channels and is half the size of the Y plane in X/Y. + src_uv.dim(0).set_bounds(0, (src_y.dim(0).extent() + 1) / 2); + src_uv.dim(1).set_bounds(0, (src_y.dim(1).extent() + 1) / 2); + src_uv.dim(2).set_bounds(0, 2); + + // Remove default memory layout constraints on the UV source so that we + // accept generic UV (including semi-planar and planar). + // + // TODO: Investigate whether it's worth specializing the cross- + // product of [semi-]planar and RGB/RGBA; this would result in 9 codepaths. + src_uv.dim(0).set_stride(Expr()); + + Halide::Func rgb_func = rgb; + Halide::OutputImageParam rgb_output = rgb_func.output_buffer(); + Halide::Expr rgb_channels = rgb_output.dim(2).extent(); + + // Specialize the generated code for RGB and RGBA. + const int vector_size = natural_vector_size(); + rgb_func.reorder(c, x, y); + rgb_func.specialize(rgb_channels == 3).unroll(c).vectorize(x, vector_size); + rgb_func.specialize(rgb_channels == 4).unroll(c).vectorize(x, vector_size); + + // Require that the output buffer be interleaved and tightly-packed; + // that is, either RGBRGBRGB[...] or RGBARGBARGBA[...], without gaps + // between pixels. + rgb_output.dim(0).set_stride(rgb_channels); + rgb_output.dim(2).set_stride(1); + + // RGB output starts at index zero in every dimension. + rgb_output.dim(0).set_min(0); + rgb_output.dim(1).set_min(0); + rgb_output.dim(2).set_min(0); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(YuvRgb, yuv_rgb_generator) diff --git a/mediapipe/util/frame_buffer/halide/yuv_rotate_generator.cc b/mediapipe/util/frame_buffer/halide/yuv_rotate_generator.cc new file mode 100644 index 000000000..46f654a23 --- /dev/null +++ b/mediapipe/util/frame_buffer/halide/yuv_rotate_generator.cc @@ -0,0 +1,91 @@ +// 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" +#include "mediapipe/util/frame_buffer/halide/common.h" + +namespace { + +using ::mediapipe::frame_buffer::halide::common::rotate; + +class YuvRotate : public Halide::Generator { + public: + Var x{"x"}, y{"y"}; + + // Input because that allows us to apply constraints on stride, etc. + Input> src_y{"src_y"}; + Input> src_uv{"src_uv"}; + // Rotation angle in degrees counter-clockwise. Must be in {0, 90, 180, 270}. + Input rotation_angle{"rotation_angle", 0}; + + Output dst_y{"dst_y", UInt(8), 2}; + Output dst_uv{"dst_uv", UInt(8), 3}; + + void generate(); + void schedule(); +}; + +void YuvRotate::generate() { + const Halide::Expr width = src_y.dim(0).extent(); + const Halide::Expr height = src_y.dim(1).extent(); + + // Rotate each of the YUV planes independently. + rotate(src_y, dst_y, width, height, rotation_angle); + rotate(src_uv, dst_uv, (width + 1) / 2, (height + 1) / 2, rotation_angle); +} + +void YuvRotate::schedule() { + // TODO: Remove specialization for (angle == 0) since that is + // a no-op and callers should simply skip rotation. Doing so would cause + // a bounds assertion crash if called with angle=0, however. + Halide::Func dst_y_func = dst_y; + dst_y_func.specialize(rotation_angle == 0).reorder(x, y); + dst_y_func.specialize(rotation_angle == 90).reorder(y, x); + dst_y_func.specialize(rotation_angle == 180).reorder(x, y); + dst_y_func.specialize(rotation_angle == 270).reorder(y, x); + + Halide::Func dst_uv_func = dst_uv; + Halide::Var c = dst_uv_func.args()[2]; + dst_uv_func.unroll(c); + dst_uv_func.specialize(rotation_angle == 0).reorder(c, x, y); + dst_uv_func.specialize(rotation_angle == 90).reorder(c, y, x); + dst_uv_func.specialize(rotation_angle == 180).reorder(c, x, y); + dst_uv_func.specialize(rotation_angle == 270).reorder(c, y, x); + + // Y plane dimensions start at zero. We could additionally constrain the + // extent to be even, but that doesn't seem to have any benefit. + Halide::OutputImageParam dst_y_output = dst_y_func.output_buffer(); + src_y.dim(0).set_min(0); + src_y.dim(1).set_min(0); + dst_y_output.dim(0).set_min(0); + dst_y_output.dim(1).set_min(0); + + // UV plane has two channels and is half the size of the Y plane in X/Y. + Halide::OutputImageParam dst_uv_output = dst_uv_func.output_buffer(); + src_uv.dim(0).set_bounds(0, (src_y.dim(0).extent() + 1) / 2); + src_uv.dim(1).set_bounds(0, (src_y.dim(1).extent() + 1) / 2); + src_uv.dim(2).set_bounds(0, 2); + dst_uv_output.dim(0).set_bounds(0, (dst_y_output.dim(0).extent() + 1) / 2); + dst_uv_output.dim(1).set_bounds(0, (dst_y_output.dim(1).extent() + 1) / 2); + dst_uv_output.dim(2).set_bounds(0, 2); + + // Remove default memory layout constraints and accept/produce generic UV + // (including semi-planar and planar). + src_uv.dim(0).set_stride(Expr()); + dst_uv_output.dim(0).set_stride(Expr()); +} + +} // namespace + +HALIDE_REGISTER_GENERATOR(YuvRotate, yuv_rotate_generator) diff --git a/mediapipe/util/frame_buffer/rgb_buffer.cc b/mediapipe/util/frame_buffer/rgb_buffer.cc new file mode 100644 index 000000000..9ae849eab --- /dev/null +++ b/mediapipe/util/frame_buffer/rgb_buffer.cc @@ -0,0 +1,132 @@ +// 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/rgb_buffer.h" + +#include + +#include "mediapipe/util/frame_buffer/buffer_common.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_gray_halide.h" +#include "mediapipe/util/frame_buffer/halide/rgb_resize_halide.h" +#include "mediapipe/util/frame_buffer/halide/rgb_rgb_halide.h" +#include "mediapipe/util/frame_buffer/halide/rgb_rotate_halide.h" +#include "mediapipe/util/frame_buffer/halide/rgb_yuv_halide.h" +#include "mediapipe/util/frame_buffer/yuv_buffer.h" + +namespace mediapipe { +namespace frame_buffer { + +RgbBuffer::RgbBuffer(uint8_t* data, int width, int height, bool alpha) + : owned_buffer_(nullptr) { + Initialize(data, width, height, alpha); +} + +RgbBuffer::RgbBuffer(uint8_t* data, int width, int height, int row_stride, + bool alpha) { + const int channels = alpha ? 4 : 3; + const halide_dimension_t dimensions[3] = {{/*m=*/0, width, channels}, + {/*m=*/0, height, row_stride}, + {/*m=*/0, channels, 1}}; + buffer_ = Halide::Runtime::Buffer(data, /*d=*/3, dimensions); +} + +RgbBuffer::RgbBuffer(int width, int height, bool alpha) + : owned_buffer_(new uint8_t[ByteSize(width, height, alpha)]) { + Initialize(owned_buffer_.get(), width, height, alpha); +} + +RgbBuffer::RgbBuffer(const RgbBuffer& other) : buffer_(other.buffer_) { + // Never copy owned_buffer; ownership remains with the source of the copy. +} + +RgbBuffer::RgbBuffer(RgbBuffer&& other) { *this = std::move(other); } + +RgbBuffer& RgbBuffer::operator=(const RgbBuffer& other) { + if (this != &other) { + buffer_ = other.buffer_; + } + return *this; +} +RgbBuffer& RgbBuffer::operator=(RgbBuffer&& other) { + if (this != &other) { + owned_buffer_ = std::move(other.owned_buffer_); + buffer_ = other.buffer_; + } + return *this; +} + +RgbBuffer::~RgbBuffer() {} + +bool RgbBuffer::Crop(int x0, int y0, int x1, int y1) { + // Twiddle the buffer start and extents to crop images. + return common::crop_buffer(x0, y0, x1, y1, buffer()); +} + +bool RgbBuffer::Resize(RgbBuffer* output) { + if (output->channels() > channels()) { + // Fail fast; the Halide implementation would otherwise output garbage + // alpha values (i.e. duplicate the blue channel into alpha). + return false; + } + const int result = rgb_resize_halide( + buffer(), static_cast(width()) / output->width(), + static_cast(height()) / output->height(), output->buffer()); + return result == 0; +} + +bool RgbBuffer::Rotate(int angle, RgbBuffer* output) { + const int result = rgb_rotate_halide(buffer(), angle, output->buffer()); + return result == 0; +} + +bool RgbBuffer::FlipHorizontally(RgbBuffer* output) { + const int result = rgb_flip_halide(buffer(), + false, // horizontal + output->buffer()); + return result == 0; +} + +bool RgbBuffer::FlipVertically(RgbBuffer* output) { + const int result = rgb_flip_halide(buffer(), + true, // vertical + output->buffer()); + return result == 0; +} + +bool RgbBuffer::Convert(YuvBuffer* output) { + const int result = + rgb_yuv_halide(buffer(), output->y_buffer(), output->uv_buffer()); + return result == 0; +} + +bool RgbBuffer::Convert(GrayBuffer* output) { + const int result = rgb_gray_halide(buffer(), output->buffer()); + return result == 0; +} + +bool RgbBuffer::Convert(RgbBuffer* output) { + const int result = rgb_rgb_halide(buffer(), 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( + data, width, height, channels); +} + +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/rgb_buffer.h b/mediapipe/util/frame_buffer/rgb_buffer.h new file mode 100644 index 000000000..06423c3f6 --- /dev/null +++ b/mediapipe/util/frame_buffer/rgb_buffer.h @@ -0,0 +1,139 @@ +// 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_RGB_BUFFER_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_RGB_BUFFER_H_ + +#include + +#include "HalideBuffer.h" +#include "HalideRuntime.h" +#include "mediapipe/util/frame_buffer/gray_buffer.h" +#include "mediapipe/util/frame_buffer/yuv_buffer.h" + +namespace mediapipe { +namespace frame_buffer { + +// RgbBuffer represents a view over an interleaved RGB/RGBA image. +// +// RgbBuffers may be copied and moved efficiently; their backing buffers are +// shared and never deep copied. +// +// RgbBuffer requires a minimum image width depending on the natural vector +// size of the platform, e.g., 16px. This is not validated by RgbBuffer. +class RgbBuffer { + public: + // Returns the size (in bytes) of an RGB/RGBA image of the given dimensions + // without padding. + static int ByteSize(int width, int height, bool alpha) { + return width * height * (alpha ? 4 : 3); + } + + // Builds a RgbBuffer using the given backing buffer and dimensions. + // + // Does not take ownership of the backing buffer (provided in 'data'). + RgbBuffer(uint8_t* data, int width, int height, bool alpha); + + // Builds a RgbBuffer using the given backing buffer and dimensions. + // 'row_stride' must be greater than or equal to 'width'. Padding bytes are at + // the end of each row, following the image bytes. + // + // Does not take ownership of the backing buffer (provided in 'data'). + RgbBuffer(uint8_t* data, int width, int height, int row_stride, bool alpha); + + // Builds a RgbBuffer using the given dimensions. + // + // The underlying backing buffer is allocated and owned by this RgbBuffer. + RgbBuffer(int width, int height, bool alpha); + + // RgbBuffer is copyable. The source retains ownership of its backing buffer. + RgbBuffer(const RgbBuffer& other); + // RgbBuffer is moveable. The source loses ownership of any backing buffers. + RgbBuffer(RgbBuffer&& other); + // RgbBuffer is assignable. + RgbBuffer& operator=(const RgbBuffer& other); + RgbBuffer& operator=(RgbBuffer&& other); + + ~RgbBuffer(); + + // Performs an in-place crop. Modifies this buffer so that the new extent + // matches that of the given crop rectangle -- (x0, y0) becomes (0, 0) and + // the new width and height are x1 - x0 + 1 and y1 - y0 + 1, respectively. + bool Crop(int x0, int y0, int x1, int y1); + + // Resize this image to match the dimensions of the given output RgbBuffer + // and places the result into its backing buffer. + // + // Performs a resize with bilinear interpolation (over four source pixels). + // Resizing with an RGB source buffer and RGBA destination is currently + // unsupported. + bool Resize(RgbBuffer* output); + + // Rotate this image into the given buffer by the given angle (90, 180, 270). + // + // Rotation is specified in degrees counter-clockwise such that when rotating + // by 90 degrees, the top-right corner of the source becomes the top-left of + // the output. The output buffer must have its height and width swapped when + // rotating by 90 or 270. + // + // Any angle values other than (90, 180, 270) are invalid. + bool Rotate(int angle, RgbBuffer* output); + + // Flip this image horizontally/vertically into the given buffer. Both buffer + // dimensions and formats must match (this method does not convert RGB-to-RGBA + // nor RGBA-to-RGB). + bool FlipHorizontally(RgbBuffer* output); + bool FlipVertically(RgbBuffer* output); + + // Performs a RGB-to-YUV color format conversion and places the result + // in the given output YuvBuffer. Both buffer dimensions must match. + bool Convert(YuvBuffer* output); + + // Performs a RGB to grayscale format conversion. + bool Convert(GrayBuffer* output); + + // Performs a rgb to rgba / rgba to rgb format conversion. + bool Convert(RgbBuffer* output); + + // Release ownership of the owned backing buffer. + uint8_t* 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. + const int width() const { return buffer_.dim(0).extent(); } + // Returns the image height. + const int height() const { return buffer_.dim(1).extent(); } + // Returns the number of color channels (3, or 4 if RGBA). + const int channels() const { return buffer_.dim(2).extent(); } + // Returns the image row stride. + const int row_stride() const { return buffer_.dim(1).stride(); } + + private: + void Initialize(uint8_t* data, int width, int height, bool alpha); + + // Non-NULL iff this RgbBuffer 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_RGB_BUFFER_H_ diff --git a/mediapipe/util/frame_buffer/rgb_buffer_test.cc b/mediapipe/util/frame_buffer/rgb_buffer_test.cc new file mode 100644 index 000000000..e5cb39c69 --- /dev/null +++ b/mediapipe/util/frame_buffer/rgb_buffer_test.cc @@ -0,0 +1,606 @@ +// 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/rgb_buffer.h" + +#include + +#include "absl/log/log.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/util/frame_buffer/gray_buffer.h" +#include "mediapipe/util/frame_buffer/yuv_buffer.h" + +// The default implementation of halide_error calls abort(), which we don't +// want. Instead, log the error and let the filter invocation fail. +extern "C" void halide_error(void*, const char* message) { + LOG(ERROR) << "Halide Error: " << message; +} + +namespace mediapipe { +namespace frame_buffer { +namespace { + +// Fill a halide_buffer_t channel with the given value. +void Fill(halide_buffer_t* buffer, int channel, int value) { + for (int y = 0; y < buffer->dim[1].extent; ++y) { + for (int x = 0; x < buffer->dim[0].extent; ++x) { + buffer->host[buffer->dim[1].stride * y + buffer->dim[0].stride * x + + buffer->dim[2].stride * channel] = value; + } + } +} + +// Fill an RgbBuffer with (0, 0, 0). Fills the alpha channel if present. +void Fill(RgbBuffer* buffer) { + for (int c = 0; c < buffer->channels(); ++c) { + Fill(buffer->buffer(), c, 0); + } +} + +// Returns a padded RGB buffer. The metadata are defined as width: 4, height: 2, +// row_stride: 18, channels: 3. +RgbBuffer GetPaddedRgbBuffer() { + static uint8_t rgb_buffer_with_padding[] = { + 10, 20, 30, 20, 30, 40, 30, 40, 50, 40, 50, 60, 0, 0, 0, 0, 0, 0, + 20, 40, 60, 40, 60, 80, 60, 80, 100, 80, 100, 120, 0, 0, 0, 0, 0, 0}; + return RgbBuffer(rgb_buffer_with_padding, + /*width=*/4, /*height=*/2, + /*row_stride=*/18, /*alpha=*/false); +} + +// Returns a padded RGB buffer. The metadata are defined as width: 4, height: 2, +// row_stride: 24, channels: 4. +RgbBuffer GetPaddedRgbaBuffer() { + static uint8_t rgb_buffer_with_padding[] = { + 10, 20, 30, 255, 20, 30, 40, 255, 30, 40, 50, 255, 40, 50, 60, 255, + 0, 0, 0, 0, 0, 0, 0, 0, 20, 40, 60, 255, 40, 60, 80, 255, + 60, 80, 100, 255, 80, 100, 120, 255, 0, 0, 0, 0, 0, 0, 0, 0}; + return RgbBuffer(rgb_buffer_with_padding, + /*width=*/4, /*height=*/2, + /*row_stride=*/24, /*alpha=*/true); +} + +// TODO: Consider move these helper methods into a util class. +// Returns true if the data in the two arrays are the same. Otherwise, return +// false. +bool CompareArray(const uint8_t* lhs_ptr, const uint8_t* rhs_ptr, int width, + int height) { + for (int i = 0; i < height; ++i) { + for (int j = 0; j < width; ++j) { + if (lhs_ptr[i * width + j] != rhs_ptr[i * width + j]) { + 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) { + if (lhs.width() != rhs.width() || lhs.height() != rhs.height()) { + return false; + } + const uint8_t* reference_ptr = const_cast(lhs).buffer()->host; + const uint8_t* converted_ptr = const_cast(rhs).buffer()->host; + return CompareArray(reference_ptr, converted_ptr, lhs.width(), lhs.height()); +} + +// Returns true if the halide buffers of two input RgbBuffer are identical. +// Otherwise, returns false; +bool CompareBuffer(const RgbBuffer& lhs, const RgbBuffer& rhs) { + if (lhs.width() != rhs.width() || lhs.height() != rhs.height() || + lhs.row_stride() != rhs.row_stride() || + lhs.channels() != rhs.channels()) { + return false; + } + const uint8_t* reference_ptr = const_cast(lhs).buffer()->host; + const uint8_t* converted_ptr = const_cast(rhs).buffer()->host; + return CompareArray(reference_ptr, converted_ptr, lhs.row_stride(), + lhs.height()); +} + +// Returns true if the halide buffers of two input YuvBuffer are identical. +// Otherwise, returns false; +bool CompareBuffer(const YuvBuffer& lhs, const YuvBuffer& rhs) { + if (lhs.width() != rhs.width() || lhs.height() != rhs.height()) { + return false; + } + const uint8_t* reference_ptr = const_cast(lhs).y_buffer()->host; + const uint8_t* converted_ptr = const_cast(rhs).y_buffer()->host; + if (!CompareArray(reference_ptr, converted_ptr, lhs.width(), lhs.height())) { + return false; + } + reference_ptr = const_cast(lhs).uv_buffer()->host; + converted_ptr = const_cast(rhs).uv_buffer()->host; + return CompareArray(reference_ptr, converted_ptr, lhs.width(), + lhs.height() / 2); +} + +TEST(RgbBufferTest, Properties) { + RgbBuffer rgb(2, 8, false), rgba(2, 8, true); + EXPECT_EQ(2, rgb.width()); + EXPECT_EQ(8, rgb.height()); + EXPECT_EQ(3, rgb.channels()); + + EXPECT_EQ(2, rgba.width()); + EXPECT_EQ(8, rgba.height()); + EXPECT_EQ(4, rgba.channels()); +} + +TEST(RgbBufferTest, PropertiesOfPaddedRgb) { + RgbBuffer rgb_buffer = GetPaddedRgbBuffer(); + EXPECT_EQ(rgb_buffer.width(), 4); + EXPECT_EQ(rgb_buffer.height(), 2); + EXPECT_EQ(rgb_buffer.row_stride(), 18); + EXPECT_EQ(rgb_buffer.channels(), 3); +} + +TEST(RgbBufferTest, PropertiesOfPaddedRgba) { + RgbBuffer rgb_buffer = GetPaddedRgbaBuffer(); + EXPECT_EQ(rgb_buffer.width(), 4); + EXPECT_EQ(rgb_buffer.height(), 2); + EXPECT_EQ(rgb_buffer.row_stride(), 24); + EXPECT_EQ(rgb_buffer.channels(), 4); +} + +TEST(RgbBufferTest, Release) { + RgbBuffer source(8, 8, true); + delete[] source.Release(); +} + +TEST(RgbBufferTest, Assign) { + RgbBuffer source(8, 8, false); + RgbBuffer sink(nullptr, 0, 0, false); + sink = source; + EXPECT_EQ(8, sink.width()); + EXPECT_EQ(8, sink.height()); + EXPECT_EQ(3, sink.channels()); + + sink = RgbBuffer(16, 16, true); + EXPECT_EQ(16, sink.width()); + EXPECT_EQ(16, sink.height()); + EXPECT_EQ(4, sink.channels()); +} + +TEST(RgbBufferTest, MoveAssign) { + RgbBuffer source(8, 8, false); + RgbBuffer sink(nullptr, 0, 0, true); + sink = std::move(source); + EXPECT_EQ(nullptr, source.Release()); + EXPECT_EQ(8, sink.width()); + EXPECT_EQ(8, sink.height()); +} + +TEST(RgbBufferTest, MoveConstructor) { + RgbBuffer source(8, 8, false); + RgbBuffer sink(std::move(source)); + EXPECT_EQ(nullptr, source.Release()); + EXPECT_EQ(8, sink.width()); + EXPECT_EQ(8, sink.height()); +} + +TEST(RgbBufferTest, RgbCrop) { + RgbBuffer source(8, 8, false); + EXPECT_TRUE(source.Crop(2, 2, 6, 6)); +} + +TEST(RgbBufferTest, RgbaCrop) { + RgbBuffer source(8, 8, true); + EXPECT_TRUE(source.Crop(2, 2, 6, 6)); +} + +// Some operations expect images with a platform-dependent minimum width +// because their implementations are vectorized. + +TEST(RgbBufferTest, RgbResize) { + RgbBuffer source(128, 8, false); + RgbBuffer result(32, 4, false); + Fill(&source); + EXPECT_TRUE(source.Resize(&result)); + + // Test odd result sizes too. + source = RgbBuffer(64, 16, false); + result = RgbBuffer(32, 7, false); + Fill(&source); + EXPECT_TRUE(source.Resize(&result)); +} + +TEST(RgbBufferTest, RgbaResize) { + RgbBuffer source(128, 8, true); + RgbBuffer result(32, 4, true); + Fill(&source); + EXPECT_TRUE(source.Resize(&result)); + + // Test odd result sizes too. + source = RgbBuffer(64, 16, true); + result = RgbBuffer(32, 7, true); + Fill(&source); + EXPECT_TRUE(source.Resize(&result)); +} + +// Note: RGB-to-RGBA conversion currently doesn't work. +TEST(RgbBufferTest, RgbResizeDifferentFormat) { + RgbBuffer source(128, 8, false); + RgbBuffer result(16, 4, true); + Fill(&source); + EXPECT_FALSE(source.Resize(&result)); +} + +TEST(RgbBufferTest, RgbaResizeDifferentFormat) { + RgbBuffer source(128, 8, true); + RgbBuffer result(16, 4, false); + Fill(&source); + EXPECT_TRUE(source.Resize(&result)); +} + +TEST(RgbBufferTest, PaddedRgbResize) { + const int target_width = 2; + const int target_height = 1; + RgbBuffer source = GetPaddedRgbBuffer(); + RgbBuffer result(target_width, target_height, /*alpha=*/false); + + ASSERT_TRUE(source.Resize(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + EXPECT_EQ(result.channels(), 3); + EXPECT_EQ(result.row_stride(), target_width * /*pixel_stride=*/3); + + uint8_t rgb_data[] = {10, 20, 30, 30, 40, 50}; + RgbBuffer rgb_buffer = + RgbBuffer(rgb_data, target_width, target_height, /*alpha=*/false); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaResize) { + const int target_width = 2; + const int target_height = 1; + RgbBuffer source = GetPaddedRgbaBuffer(); + RgbBuffer result(target_width, target_height, /*alpha=*/true); + + ASSERT_TRUE(source.Resize(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + EXPECT_EQ(result.channels(), 4); + EXPECT_EQ(result.row_stride(), target_width * /*pixel_stride=*/4); + + uint8_t rgb_data[] = {10, 20, 30, 255, 30, 40, 50, 255}; + RgbBuffer rgb_buffer = + RgbBuffer(rgb_data, target_width, target_height, /*alpha=*/true); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, RgbRotateCheckSize) { + RgbBuffer source(4, 8, false); + RgbBuffer result(8, 4, false); + Fill(&source); + EXPECT_TRUE(source.Rotate(90, &result)); +} + +TEST(RgbBufferTest, RgbRotateCheckData) { + uint8_t* data = new uint8_t[12]; + data[0] = data[1] = data[2] = 1; // Pixel 1 + data[3] = data[4] = data[5] = 2; // Pixel 2 + data[6] = data[7] = data[8] = 3; // Pixel 3 + data[9] = data[10] = data[11] = 4; // Pixel 4 + RgbBuffer source(data, 2, 2, false); + RgbBuffer result(2, 2, false); + source.Rotate(90, &result); + EXPECT_EQ(2, result.buffer()->host[0]); + EXPECT_EQ(4, result.buffer()->host[3]); + EXPECT_EQ(1, result.buffer()->host[6]); + EXPECT_EQ(3, result.buffer()->host[9]); + delete[] data; +} + +TEST(RgbBufferTest, RgbRotateDifferentFormat) { + RgbBuffer source(4, 8, true); + RgbBuffer result(8, 4, false); + Fill(&source); + EXPECT_TRUE(source.Rotate(90, &result)); +} + +// Note: RGB-to-RGBA conversion currently doesn't work. +TEST(RgbBufferTest, RgbRotateDifferentFormatFail) { + RgbBuffer source(4, 8, false); + RgbBuffer result(8, 4, true); + Fill(&source); + EXPECT_FALSE(source.Rotate(90, &result)); +} + +TEST(RgbBufferTest, RgbaRotate) { + RgbBuffer source(4, 8, true); + RgbBuffer result(8, 4, true); + Fill(&source); + EXPECT_TRUE(source.Rotate(90, &result)); +} + +TEST(RgbBufferTest, RgbaRotateDifferentFormat) { + RgbBuffer source(4, 8, true); + RgbBuffer result(8, 4, false); + Fill(&source); + EXPECT_TRUE(source.Rotate(90, &result)); +} + +// Note: RGB-to-RGBA conversion currently doesn't work. +TEST(RgbBufferTest, RgbaRotateDifferentFormatFail) { + RgbBuffer source(4, 8, false); + RgbBuffer result(8, 4, true); + Fill(&source); + EXPECT_FALSE(source.Rotate(90, &result)); +} + +TEST(RgbBufferTest, PaddedRgbRotateCheckData) { + const int target_width = 2; + const int target_height = 4; + RgbBuffer source = GetPaddedRgbBuffer(); + RgbBuffer result(target_width, target_height, /*alpha=*/false); + + ASSERT_TRUE(source.Rotate(/*angle=*/90, &result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + EXPECT_EQ(result.channels(), 3); + EXPECT_EQ(result.row_stride(), target_width * /*pixel_stride=*/3); + + uint8_t rgb_data[] = {40, 50, 60, 80, 100, 120, 30, 40, 50, 60, 80, 100, + 20, 30, 40, 40, 60, 80, 10, 20, 30, 20, 40, 60}; + RgbBuffer rgb_buffer = + RgbBuffer(rgb_data, target_width, target_height, /*alpha=*/false); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaRotateCheckData) { + const int target_width = 2; + const int target_height = 4; + RgbBuffer result(target_width, target_height, /*alpha=*/true); + RgbBuffer source = GetPaddedRgbaBuffer(); + + ASSERT_TRUE(source.Rotate(/*angle=*/90, &result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + EXPECT_EQ(result.channels(), 4); + EXPECT_EQ(result.row_stride(), target_width * /*pixel_stride=*/4); + + uint8_t rgb_data[] = {40, 50, 60, 255, 80, 100, 120, 255, 30, 40, 50, + 255, 60, 80, 100, 255, 20, 30, 40, 255, 40, 60, + 80, 255, 10, 20, 30, 255, 20, 40, 60, 255}; + RgbBuffer rgb_buffer = + RgbBuffer(rgb_data, target_width, target_height, /*alpha=*/true); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, RgbaFlip) { + RgbBuffer source(16, 16, true); + RgbBuffer result(16, 16, true); + Fill(&source); + EXPECT_TRUE(source.FlipHorizontally(&result)); + EXPECT_TRUE(source.FlipVertically(&result)); +} + +// Note: Neither RGBA-to-RGB nor RGB-to-RGBA conversion currently works. +TEST(RgbBufferTest, RgbaFlipDifferentFormatFail) { + RgbBuffer source(16, 16, false); + RgbBuffer result(16, 16, true); + Fill(&source); + Fill(&result); + EXPECT_FALSE(source.FlipHorizontally(&result)); + EXPECT_FALSE(result.FlipHorizontally(&source)); + EXPECT_FALSE(source.FlipVertically(&result)); + EXPECT_FALSE(result.FlipVertically(&source)); +} + +TEST(RgbBufferTest, PaddedRgbFlipHorizontally) { + const int target_width = 4; + const int target_height = 2; + RgbBuffer result(target_width, target_height, /*alpha=*/false); + RgbBuffer source = GetPaddedRgbBuffer(); + + ASSERT_TRUE(source.FlipHorizontally(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + EXPECT_EQ(result.channels(), 3); + EXPECT_EQ(result.row_stride(), target_width * /*pixel_stride=*/3); + + uint8_t rgb_data[] = {40, 50, 60, 30, 40, 50, 20, 30, 40, 10, 20, 30, + 80, 100, 120, 60, 80, 100, 40, 60, 80, 20, 40, 60}; + RgbBuffer rgb_buffer = + RgbBuffer(rgb_data, target_width, target_height, /*alpha=*/false); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaFlipHorizontally) { + const int target_width = 4; + const int target_height = 2; + RgbBuffer result(target_width, target_height, /*alpha=*/true); + RgbBuffer source = GetPaddedRgbaBuffer(); + + ASSERT_TRUE(source.FlipHorizontally(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + EXPECT_EQ(result.channels(), 4); + EXPECT_EQ(result.row_stride(), target_width * /*pixel_stride=*/4); + + uint8_t rgb_data[] = {40, 50, 60, 255, 30, 40, 50, 255, 20, 30, 40, + 255, 10, 20, 30, 255, 80, 100, 120, 255, 60, 80, + 100, 255, 40, 60, 80, 255, 20, 40, 60, 255}; + RgbBuffer rgb_buffer = + RgbBuffer(rgb_data, target_width, target_height, /*alpha=*/true); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, RgbConvertNv21) { + RgbBuffer source(32, 8, false); + YuvBuffer result(32, 8, YuvBuffer::NV21); + Fill(&source); + EXPECT_TRUE(source.Convert(&result)); +} + +TEST(RgbBufferTest, RgbaConvertNv21) { + RgbBuffer source(32, 8, true); + YuvBuffer result(32, 8, YuvBuffer::NV21); + Fill(&source); + EXPECT_TRUE(source.Convert(&result)); +} + +TEST(RgbBufferTest, PaddedRgbConvertNv21) { + const int target_width = 4; + const int target_height = 2; + YuvBuffer result(target_width, target_height, YuvBuffer::NV21); + RgbBuffer source = GetPaddedRgbBuffer(); + + ASSERT_TRUE(source.Convert(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + + uint8_t yuv_data[] = {18, 28, 38, 48, 36, 56, 76, 96, 122, 135, 122, 135}; + YuvBuffer yuv_buffer = + YuvBuffer(yuv_data, target_width, target_height, YuvBuffer::NV21); + EXPECT_TRUE(CompareBuffer(yuv_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaConvertNv21) { + const int target_width = 4; + const int target_height = 2; + YuvBuffer result(target_width, target_height, YuvBuffer::NV21); + RgbBuffer source = GetPaddedRgbaBuffer(); + + ASSERT_TRUE(source.Convert(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + + uint8_t yuv_data[] = {18, 28, 38, 48, 36, 56, 76, 96, 122, 135, 122, 135}; + YuvBuffer yuv_buffer = + YuvBuffer(yuv_data, target_width, target_height, YuvBuffer::NV21); + EXPECT_TRUE(CompareBuffer(yuv_buffer, result)); +} + +TEST(RgbBufferTest, RgbConvertGray) { + uint8_t* data = new uint8_t[6]; + data[0] = 200; + data[1] = 100; + data[2] = 0; + data[3] = 0; + data[4] = 200; + data[5] = 100; + RgbBuffer source(data, 2, 1, false); + GrayBuffer result(2, 1); + EXPECT_TRUE(source.Convert(&result)); + EXPECT_EQ(118, result.buffer()->host[0]); + EXPECT_EQ(129, result.buffer()->host[1]); + delete[] data; +} + +TEST(RgbBufferTest, RgbaConvertGray) { + uint8_t* data = new uint8_t[8]; + data[0] = 200; + data[1] = 100; + data[2] = 0; + data[3] = 1; + data[4] = 0; + data[5] = 200; + data[6] = 100; + data[7] = 50; + RgbBuffer source(data, 2, 1, true); + GrayBuffer result(2, 1); + EXPECT_TRUE(source.Convert(&result)); + EXPECT_EQ(118, result.buffer()->host[0]); + EXPECT_EQ(129, result.buffer()->host[1]); + delete[] data; +} + +TEST(RgbBufferTest, PaddedRgbConvertGray) { + const int target_width = 4; + const int target_height = 2; + GrayBuffer result(target_width, target_height); + RgbBuffer source = GetPaddedRgbBuffer(); + + ASSERT_TRUE(source.Convert(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + + uint8_t gray_data[] = {18, 28, 38, 48, 36, 56, 76, 96}; + GrayBuffer gray_buffer = GrayBuffer(gray_data, target_width, target_height); + EXPECT_TRUE(CompareBuffer(gray_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaConvertGray) { + const int target_width = 4; + const int target_height = 2; + GrayBuffer result(target_width, target_height); + RgbBuffer source = GetPaddedRgbaBuffer(); + + ASSERT_TRUE(source.Convert(&result)); + EXPECT_EQ(result.width(), target_width); + EXPECT_EQ(result.height(), target_height); + + uint8_t gray_data[] = {18, 28, 38, 48, 36, 56, 76, 96}; + GrayBuffer gray_buffer = GrayBuffer(gray_data, target_width, target_height); + EXPECT_TRUE(CompareBuffer(gray_buffer, result)); +} + +TEST(RgbBufferTest, RgbConvertRgba) { + constexpr int kWidth = 2, kHeight = 1; + uint8_t rgb_data[] = {200, 100, 50, 100, 50, 20}; + RgbBuffer source(rgb_data, kWidth, kHeight, false); + RgbBuffer result(kWidth, kHeight, true); + + ASSERT_TRUE(source.Convert(&result)); + + uint8_t rgba_data[] = {200, 100, 50, 255, 100, 50, 20, 255}; + RgbBuffer rgba_buffer = RgbBuffer(rgba_data, kWidth, kHeight, true); + EXPECT_TRUE(CompareBuffer(rgba_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbConvertRgba) { + constexpr int kWidth = 4, kHeight = 2; + RgbBuffer source = GetPaddedRgbBuffer(); + RgbBuffer result(kWidth, kHeight, true); + ASSERT_TRUE(source.Convert(&result)); + + uint8_t rgba_data[]{10, 20, 30, 255, 20, 30, 40, 255, 30, 40, 50, + 255, 40, 50, 60, 255, 20, 40, 60, 255, 40, 60, + 80, 255, 60, 80, 100, 255, 80, 100, 120, 255}; + RgbBuffer rgba_buffer = RgbBuffer(rgba_data, kWidth, kHeight, true); + EXPECT_TRUE(CompareBuffer(rgba_buffer, result)); +} + +TEST(RgbBufferTest, RgbaConvertRgb) { + constexpr int kWidth = 2, kHeight = 1; + uint8_t rgba_data[] = {200, 100, 50, 30, 100, 50, 20, 70}; + RgbBuffer source(rgba_data, kWidth, kHeight, true); + RgbBuffer result(kWidth, kHeight, false); + + ASSERT_TRUE(source.Convert(&result)); + + uint8_t rgb_data[] = {200, 100, 50, 100, 50, 20}; + RgbBuffer rgb_buffer = RgbBuffer(rgb_data, kWidth, kHeight, false); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} + +TEST(RgbBufferTest, PaddedRgbaConvertRgb) { + constexpr int kWidth = 4, kHeight = 2; + RgbBuffer source = GetPaddedRgbaBuffer(); + RgbBuffer result(kWidth, kHeight, false); + + ASSERT_TRUE(source.Convert(&result)); + + uint8_t rgb_data[] = {10, 20, 30, 20, 30, 40, 30, 40, 50, 40, 50, 60, + 20, 40, 60, 40, 60, 80, 60, 80, 100, 80, 100, 120}; + RgbBuffer rgb_buffer = RgbBuffer(rgb_data, kWidth, kHeight, false); + EXPECT_TRUE(CompareBuffer(rgb_buffer, result)); +} +} // namespace +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/yuv_buffer.cc b/mediapipe/util/frame_buffer/yuv_buffer.cc new file mode 100644 index 000000000..f96282134 --- /dev/null +++ b/mediapipe/util/frame_buffer/yuv_buffer.cc @@ -0,0 +1,152 @@ +// 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/yuv_buffer.h" + +#include + +#include "mediapipe/util/frame_buffer/buffer_common.h" +#include "mediapipe/util/frame_buffer/halide/yuv_flip_halide.h" +#include "mediapipe/util/frame_buffer/halide/yuv_resize_halide.h" +#include "mediapipe/util/frame_buffer/halide/yuv_rgb_halide.h" +#include "mediapipe/util/frame_buffer/halide/yuv_rotate_halide.h" +#include "mediapipe/util/frame_buffer/rgb_buffer.h" + +namespace mediapipe { +namespace frame_buffer { + +YuvBuffer::YuvBuffer(uint8_t* y_plane, uint8_t* u_plane, uint8_t* v_plane, + int width, int height, int row_stride_y, int row_stride_uv, + int pixel_stride_uv) { + // Initialize the buffer shapes: {min, extent, stride} per dimension. + // TODO: Ensure that width is less than or equal to row stride. + const halide_dimension_t y_dimensions[2] = { + {0, width, 1}, + {0, height, row_stride_y}, + }; + y_buffer_ = Halide::Runtime::Buffer(y_plane, 2, y_dimensions); + + // Note that the Halide implementation expects the planes to be in VU + // order, so we point at the V plane first. + const halide_dimension_t uv_dimensions[3] = { + {0, (width + 1) / 2, pixel_stride_uv}, + {0, (height + 1) / 2, row_stride_uv}, + {0, 2, static_cast(u_plane - v_plane)}, + }; + uv_buffer_ = Halide::Runtime::Buffer(v_plane, 3, uv_dimensions); +} + +YuvBuffer::YuvBuffer(uint8_t* data, int width, int height, Format format) + : owned_buffer_(nullptr) { + Initialize(data, width, height, format); +} + +YuvBuffer::YuvBuffer(int width, int height, Format format) + : owned_buffer_(new uint8_t[ByteSize(width, height)]) { + Initialize(owned_buffer_.get(), width, height, format); +} + +YuvBuffer::YuvBuffer(const YuvBuffer& other) + : y_buffer_(other.y_buffer_), uv_buffer_(other.uv_buffer_) { + // Never copy owned_buffer; ownership remains with the source of the copy. +} + +YuvBuffer::YuvBuffer(YuvBuffer&& other) { *this = std::move(other); } + +YuvBuffer& YuvBuffer::operator=(const YuvBuffer& other) { + if (this != &other) { + y_buffer_ = other.y_buffer_; + uv_buffer_ = other.uv_buffer_; + } + return *this; +} +YuvBuffer& YuvBuffer::operator=(YuvBuffer&& other) { + if (this != &other) { + owned_buffer_ = std::move(other.owned_buffer_); + y_buffer_ = other.y_buffer_; + uv_buffer_ = other.uv_buffer_; + } + return *this; +} + +YuvBuffer::~YuvBuffer() {} + +void YuvBuffer::Initialize(uint8_t* data, int width, int height, + Format format) { + y_buffer_ = Halide::Runtime::Buffer(data, width, height); + + uint8_t* uv_data = data + (width * height); + switch (format) { + case NV21: + // Interleaved UV (actually VU order). + uv_buffer_ = Halide::Runtime::Buffer::make_interleaved( + uv_data, (width + 1) / 2, (height + 1) / 2, 2); + break; + case YV12: + // Planar UV (actually VU order). + uv_buffer_ = Halide::Runtime::Buffer(uv_data, (width + 1) / 2, + (height + 1) / 2, 2); + // NOTE: Halide operations have not been tested extensively in this + // configuration. + break; + } +} + +bool YuvBuffer::Crop(int x0, int y0, int x1, int y1) { + if (x0 & 1 || y0 & 1) { + // YUV images must be left-and top-aligned to even X/Y coordinates. + return false; + } + + // Twiddle the buffer start and extents for each plane to crop images. + return (common::crop_buffer(x0, y0, x1, y1, y_buffer()) && + common::crop_buffer(x0 / 2, y0 / 2, x1 / 2, y1 / 2, uv_buffer())); +} + +bool YuvBuffer::Resize(YuvBuffer* output) { + const int result = yuv_resize_halide( + y_buffer(), uv_buffer(), static_cast(width()) / output->width(), + static_cast(height()) / output->height(), output->y_buffer(), + output->uv_buffer()); + return result == 0; +} + +bool YuvBuffer::Rotate(int angle, YuvBuffer* output) { + const int result = yuv_rotate_halide(y_buffer(), uv_buffer(), angle, + output->y_buffer(), output->uv_buffer()); + return result == 0; +} + +bool YuvBuffer::FlipHorizontally(YuvBuffer* output) { + const int result = yuv_flip_halide(y_buffer(), uv_buffer(), + false, // horizontal + output->y_buffer(), output->uv_buffer()); + return result == 0; +} + +bool YuvBuffer::FlipVertically(YuvBuffer* output) { + const int result = yuv_flip_halide(y_buffer(), uv_buffer(), + true, // vertical + output->y_buffer(), output->uv_buffer()); + return result == 0; +} + +bool YuvBuffer::Convert(bool halve, RgbBuffer* output) { + const int result = + yuv_rgb_halide(y_buffer(), uv_buffer(), halve, output->buffer()); + return result == 0; +} + +} // namespace frame_buffer +} // namespace mediapipe diff --git a/mediapipe/util/frame_buffer/yuv_buffer.h b/mediapipe/util/frame_buffer/yuv_buffer.h new file mode 100644 index 000000000..8e4715dcf --- /dev/null +++ b/mediapipe/util/frame_buffer/yuv_buffer.h @@ -0,0 +1,160 @@ +// 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_YUV_BUFFER_H_ +#define MEDIAPIPE_UTIL_FRAME_BUFFER_YUV_BUFFER_H_ + +#include + +#include "HalideBuffer.h" +#include "HalideRuntime.h" + +namespace mediapipe { +namespace frame_buffer { +class RgbBuffer; + +// YuvBuffer represents a view over a YUV 4:2:0 image. +// +// YuvBuffers may be copied and moved efficiently; their backing buffers are +// shared and never deep copied. +// +// YuvBuffer requires a minimum image width depending on the natural vector +// size of the platform, e.g., 16px. This is not validated by YuvBuffer. +class YuvBuffer { + public: + // YUV formats. Rather than supporting every possible format, we prioritize + // formats with broad hardware/platform support. + // + // Enum values are FourCC codes; see http://fourcc.org/yuv.php for more. + enum Format { + NV21 = 0x3132564E, // YUV420SP (VU interleaved) + YV12 = 0x32315659, // YUV420P (VU planar) + }; + + // Returns the size (in bytes) of a YUV image of the given dimensions. + static int ByteSize(int width, int height) { + // 1 byte per pixel in the Y plane, 2 bytes per 2x2 block in the UV plane. + // Dimensions with odd sizes are rounded up. + const int y_size = width * height; + const int uv_size = ((width + 1) / 2) * ((height + 1) / 2) * 2; + return y_size + uv_size; + } + + // Builds a generic YUV420 YuvBuffer with the given backing buffers, + // dimensions and strides. Supports both interleaved or planar UV with + // custom strides. + // + // Does not take ownership of any backing buffers, which must be large + // enough to fit their contents. + YuvBuffer(uint8_t* y_plane, uint8_t* u_plane, uint8_t* v_plane, int width, + int height, int row_stride_y, int row_stride_uv, + int pixel_stride_uv); + + // Builds a YuvBuffer using the given backing buffer, dimensions, and format. + // Expects an NV21- or YV12-format image only. + // + // Does not take ownership of the backing buffer (provided in 'data'), which + // must be sized to hold at least the amount indicated by ByteSize(). + YuvBuffer(uint8_t* data, int width, int height, Format format); + + // Builds a YuvBuffer using the given dimensions and format. Expects + // an NV21- or YV12-format image only. + // + // The underlying backing buffer is allocated and owned by this YuvBuffer. + YuvBuffer(int width, int height, Format format); + + // YuvBuffer is copyable. The source retains ownership of its backing buffers. + YuvBuffer(const YuvBuffer& other); + // YuvBuffer is moveable. The source loses ownership of any backing buffers. + YuvBuffer(YuvBuffer&& other); + // YuvBuffer is assignable. + YuvBuffer& operator=(const YuvBuffer& other); + YuvBuffer& operator=(YuvBuffer&& other); + + ~YuvBuffer(); + + // Performs an in-place crop. Modifies this buffer so that the new extent + // matches that of the given crop rectangle -- (x0, y0) becomes (0, 0) and + // the new width and height are x1 - x0 + 1 and y1 - y0 + 1, respectively. + // + // Note that the top-left corner (x0, y0) coordinates must be even to + // maintain alignment between the Y and UV grids. + bool Crop(int x0, int y0, int x1, int y1); + + // Resize this image to match the dimensions of the given output YuvBuffer + // and places the result into its backing buffer. + // + // Performs a resize with bilinear interpolation (over four source pixels). + bool Resize(YuvBuffer* output); + + // Rotate this image into the given buffer by the given angle (90, 180, 270). + // + // Rotation is specified in degrees counter-clockwise such that when rotating + // by 90 degrees, the top-right corner of the source becomes the top-left of + // the output. The output buffer must have its height and width swapped when + // rotating by 90 or 270. + // + // Any angle values other than (90, 180, 270) are invalid. + bool Rotate(int angle, YuvBuffer* output); + + // Flip this image horizontally/vertically into the given buffer. Both buffer + // dimensions must match. + bool FlipHorizontally(YuvBuffer* output); + bool FlipVertically(YuvBuffer* output); + + // Performs a YUV-to-RGB color format conversion and places the result + // in the given output RgbBuffer. Both buffer dimensions must match. + // + // When halve is true, the converted output is downsampled by a factor of + // two by discarding three of four luminance values in every 2x2 block. + bool Convert(bool halve, RgbBuffer* output); + + // Release ownership of the owned backing buffer. + uint8_t* Release() { return owned_buffer_.release(); } + + // Returns the halide_buffer_t* for the Y plane. + const halide_buffer_t* y_buffer() const { return y_buffer_.raw_buffer(); } + // Returns the halide_buffer_t* for the UV plane(s). + const halide_buffer_t* uv_buffer() const { return uv_buffer_.raw_buffer(); } + // Returns the halide_buffer_t* for the Y plane. + halide_buffer_t* y_buffer() { return y_buffer_.raw_buffer(); } + // Returns the halide_buffer_t* for the UV plane(s). + halide_buffer_t* uv_buffer() { return uv_buffer_.raw_buffer(); } + + // Returns the image width. + const int width() const { return y_buffer_.dim(0).extent(); } + // Returns the image height. + const int height() const { return y_buffer_.dim(1).extent(); } + + private: + void Initialize(uint8_t* data, int width, int height, Format format); + + // Non-NULL iff this YuvBuffer owns its buffer. + std::unique_ptr owned_buffer_; + + // Y (luminance) backing buffer: layout is always width x height. + Halide::Runtime::Buffer y_buffer_; + + // UV (chrominance) backing buffer; width/2 x height/2 x 2 (channel). + // May be interleaved or planar. + // + // Note that the planes are in the reverse of the usual order: channel 0 is V + // and channel 1 is U. + Halide::Runtime::Buffer uv_buffer_; +}; + +} // namespace frame_buffer +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_YUV_BUFFER_H_ diff --git a/mediapipe/util/frame_buffer/yuv_buffer_test.cc b/mediapipe/util/frame_buffer/yuv_buffer_test.cc new file mode 100644 index 000000000..a18b19a92 --- /dev/null +++ b/mediapipe/util/frame_buffer/yuv_buffer_test.cc @@ -0,0 +1,251 @@ +// 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/yuv_buffer.h" + +#include + +#include "absl/log/log.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/util/frame_buffer/rgb_buffer.h" + +// The default implementation of halide_error calls abort(), which we don't +// want. Instead, log the error and let the filter invocation fail. +extern "C" void halide_error(void*, const char* message) { + LOG(ERROR) << "Halide Error: " << message; +} + +namespace mediapipe { +namespace frame_buffer { +namespace { + +// Fill a halide_buffer_t channel with the given value. +void Fill(halide_buffer_t* buffer, int channel, int value) { + for (int y = 0; y < buffer->dim[1].extent; ++y) { + for (int x = 0; x < buffer->dim[0].extent; ++x) { + buffer->host[buffer->dim[1].stride * y + buffer->dim[0].stride * x + + buffer->dim[2].stride * channel] = value; + } + } +} + +// Fill a YuvBuffer with the given YUV color. +void Fill(YuvBuffer* buffer, uint8_t y, uint8_t u, uint8_t v) { + Fill(buffer->y_buffer(), 0, y); + Fill(buffer->uv_buffer(), 1, u); + Fill(buffer->uv_buffer(), 0, v); +} + +TEST(YuvBufferTest, Properties) { + YuvBuffer yuv(2, 8, YuvBuffer::NV21); + EXPECT_EQ(2, yuv.width()); + EXPECT_EQ(8, yuv.height()); +} + +TEST(YuvBufferTest, Release) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + delete[] source.Release(); +} + +TEST(YuvBufferTest, Assign) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + YuvBuffer sink(nullptr, 0, 0, YuvBuffer::NV21); + sink = source; + EXPECT_EQ(8, sink.width()); + EXPECT_EQ(8, sink.height()); + + sink = YuvBuffer(16, 16, YuvBuffer::NV21); + EXPECT_EQ(16, sink.width()); + EXPECT_EQ(16, sink.height()); +} + +TEST(YuvBufferTest, MoveAssign) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + YuvBuffer sink(nullptr, 0, 0, YuvBuffer::NV21); + sink = std::move(source); + EXPECT_EQ(nullptr, source.Release()); + EXPECT_EQ(8, sink.width()); + EXPECT_EQ(8, sink.height()); +} + +TEST(YuvBufferTest, MoveConstructor) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + YuvBuffer sink(std::move(source)); + EXPECT_EQ(nullptr, source.Release()); + EXPECT_EQ(8, sink.width()); + EXPECT_EQ(8, sink.height()); +} + +TEST(YuvBufferTest, GenericSemiplanarLayout) { + uint8_t y_plane[16], uv_plane[8]; + YuvBuffer buffer(y_plane, uv_plane, uv_plane + 1, 4, 4, 4, 4, 2); + Fill(&buffer, 16, 32, 64); + + for (int i = 0; i < 16; ++i) { + EXPECT_EQ(y_plane[i], 16) << i; + } + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(uv_plane[2 * i], 32); + EXPECT_EQ(uv_plane[2 * i + 1], 64); + } +} + +TEST(YuvBufferTest, GenericPlanarLayout) { + uint8_t y_plane[16], u_plane[4], v_plane[4]; + YuvBuffer buffer(y_plane, u_plane, v_plane, 4, 4, 4, 2, 1); + Fill(&buffer, 16, 32, 64); + + for (int i = 0; i < 16; ++i) { + EXPECT_EQ(y_plane[i], 16) << i; + } + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(u_plane[i], 32); + EXPECT_EQ(v_plane[i], 64); + } +} + +TEST(YuvBufferTest, Nv21Crop) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + EXPECT_TRUE(source.Crop(2, 2, 6, 6)); +} + +TEST(YuvBufferTest, Nv21Resize) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + YuvBuffer result(4, 4, YuvBuffer::NV21); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.Resize(&result)); + + // Test odd result sizes too. + source = YuvBuffer(500, 362, YuvBuffer::NV21); + result = YuvBuffer(320, 231, YuvBuffer::NV21); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.Resize(&result)); +} + +TEST(YuvBufferTest, Nv21ResizeDifferentFormat) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + YuvBuffer result(4, 4, YuvBuffer::YV12); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.Resize(&result)); +} + +TEST(YuvBufferTest, Nv21Rotate) { + YuvBuffer source(4, 8, YuvBuffer::NV21); + YuvBuffer result(8, 4, YuvBuffer::NV21); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.Rotate(90, &result)); +} + +TEST(YuvBufferTest, Nv21RotateDifferentFormat) { + YuvBuffer source(8, 8, YuvBuffer::NV21); + YuvBuffer result(8, 8, YuvBuffer::YV12); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.Rotate(90, &result)); +} + +TEST(YuvBufferTest, Nv21RotateFailBounds) { + // Expect failure if the destination doesn't have the correct bounds. + YuvBuffer source(4, 8, YuvBuffer::NV21); + YuvBuffer result(4, 8, YuvBuffer::NV21); + Fill(&source, 16, 32, 64); + EXPECT_FALSE(source.Rotate(90, &result)); +} + +TEST(YuvBufferTest, Nv21Flip) { + YuvBuffer source(16, 16, YuvBuffer::NV21); + YuvBuffer result(16, 16, YuvBuffer::NV21); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.FlipHorizontally(&result)); + EXPECT_TRUE(source.FlipVertically(&result)); +} + +TEST(YuvBufferTest, Nv21FlipDifferentFormat) { + YuvBuffer source(16, 16, YuvBuffer::NV21); + YuvBuffer result(16, 16, YuvBuffer::YV12); + Fill(&source, 16, 32, 64); + EXPECT_TRUE(source.FlipHorizontally(&result)); + EXPECT_TRUE(source.FlipVertically(&result)); +} + +TEST(YuvBufferTest, Nv21ConvertRgb) { + // Note that RGB conversion expects at least images of width >= 32 because + // the implementation is vectorized. + YuvBuffer source(32, 8, YuvBuffer::NV21); + Fill(&source, 52, 170, 90); + + RgbBuffer result_rgb(32, 8, false); + EXPECT_TRUE(source.Convert(false, &result_rgb)); + + RgbBuffer result_rgba(32, 8, true); + EXPECT_TRUE(source.Convert(false, &result_rgba)); + + uint8_t* pixels = result_rgba.buffer()->host; + ASSERT_TRUE(pixels); + EXPECT_EQ(pixels[0], 0); + EXPECT_EQ(pixels[1], 65); + EXPECT_EQ(pixels[2], 126); + EXPECT_EQ(pixels[3], 255); +} + +TEST(YuvBufferTest, Nv21ConvertRgbCropped) { + // Note that RGB conversion expects at least images of width >= 32 because + // the implementation is vectorized. + YuvBuffer source(1024, 768, YuvBuffer::NV21); + Fill(&source, 52, 170, 90); + + // YUV images must be left-and top-aligned to even X/Y coordinates, + // regardless of whether the target image has even or odd width/height. + EXPECT_FALSE(source.Crop(1, 1, 512, 384)); + EXPECT_FALSE(source.Crop(1, 1, 511, 383)); + + YuvBuffer source1(source); + EXPECT_TRUE(source1.Crop(64, 64, 512, 384)); + RgbBuffer result_rgb(source1.width(), source1.height(), false); + EXPECT_TRUE(source1.Convert(false, &result_rgb)); + + YuvBuffer source2(source); + EXPECT_TRUE(source2.Crop(64, 64, 511, 383)); + RgbBuffer result_rgba(source2.width(), source2.height(), true); + EXPECT_TRUE(source2.Convert(false, &result_rgba)); + + uint8_t* pixels = result_rgba.buffer()->host; + ASSERT_TRUE(pixels); + EXPECT_EQ(pixels[0], 0); + EXPECT_EQ(pixels[1], 65); + EXPECT_EQ(pixels[2], 126); + EXPECT_EQ(pixels[3], 255); +} + +TEST(YuvBufferTest, Nv21ConvertRgbHalve) { + YuvBuffer source(64, 8, YuvBuffer::NV21); + Fill(&source, 52, 170, 90); + + RgbBuffer result_rgb(32, 4, false); + EXPECT_TRUE(source.Convert(true, &result_rgb)); + + RgbBuffer result_rgba(32, 4, true); + EXPECT_TRUE(source.Convert(true, &result_rgba)); + + uint8_t* pixels = result_rgba.buffer()->host; + ASSERT_TRUE(pixels); + EXPECT_EQ(pixels[0], 0); + EXPECT_EQ(pixels[1], 65); + EXPECT_EQ(pixels[2], 126); + EXPECT_EQ(pixels[3], 255); +} + +} // namespace +} // namespace frame_buffer +} // namespace mediapipe diff --git a/third_party/halide.BUILD b/third_party/halide.BUILD new file mode 100644 index 000000000..5578c5650 --- /dev/null +++ b/third_party/halide.BUILD @@ -0,0 +1,70 @@ +# 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. + +load("@halide//:halide.bzl", "halide_language_copts") + +licenses(["notice"]) + +package( + default_visibility = ["//visibility:public"], +) + +cc_library( + name = "language", + hdrs = ["include/Halide.h"], + copts = halide_language_copts(), + includes = ["include"], + deps = [ + ":runtime", + ], +) + +cc_library( + name = "runtime", + hdrs = glob([ + "include/HalideRuntime*.h", + "include/HalideBuffer*.h", + ]), + includes = ["include"], +) + +cc_library( + name = "lib_halide_static", + srcs = select({ + "@mediapipe//mediapipe:windows": [ + "lib/Release/Halide.lib", + "bin/Release/Halide.dll", + ], + "//conditions:default": [ + "lib/libHalide.a", + ], + }), + visibility = ["//visibility:private"], +) + +cc_library( + name = "gengen", + srcs = [ + "share/Halide/tools/GenGen.cpp", + ], + includes = [ + "include", + "share/Halide/tools", + ], + visibility = ["//visibility:public"], + deps = [ + ":language", + ":lib_halide_static", + ], +) diff --git a/third_party/halide/BUILD b/third_party/halide/BUILD new file mode 100644 index 000000000..82bab3ffd --- /dev/null +++ b/third_party/halide/BUILD @@ -0,0 +1 @@ +# This empty BUILD file is required to make Bazel treat this directory as a package. diff --git a/third_party/halide/BUILD.bazel b/third_party/halide/BUILD.bazel new file mode 100644 index 000000000..cda994204 --- /dev/null +++ b/third_party/halide/BUILD.bazel @@ -0,0 +1,40 @@ +# 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. + +load("@halide//:halide.bzl", "halide_library_runtimes") + +licenses(["notice"]) + +package( + default_visibility = ["//visibility:public"], +) + +halide_library_runtimes() + +# Aliases to platform-specific targets. +[ + alias( + name = target_name, + actual = select({ + "//conditions:default": "@linux_halide//:" + target_name, + "@mediapipe//mediapipe:macos": "@macos_halide//:" + target_name, + "@mediapipe//mediapipe:windows": "@windows_halide//:" + target_name, + }), + ) + for target_name in [ + "language", + "runtime", + "gengen", + ] +] diff --git a/third_party/halide/halide.bzl b/third_party/halide/halide.bzl new file mode 100644 index 000000000..68d31fe48 --- /dev/null +++ b/third_party/halide/halide.bzl @@ -0,0 +1,875 @@ +# Copyright 2023 The MediaPipe Authors. All rights reserved. +# +# 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. + +"""Bazel build rules for Halide.""" + +load("@bazel_skylib//lib:collections.bzl", "collections") +load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "use_cpp_toolchain") + +def halide_language_copts(): + _common_opts = [ + "-fPIC", + "-frtti", + "-Wno-conversion", + "-Wno-sign-compare", + ] + _posix_opts = [ + "$(STACK_FRAME_UNLIMITED)", + "-fno-exceptions", + "-funwind-tables", + "-fvisibility-inlines-hidden", + ] + _msvc_opts = [ + "-D_CRT_SECURE_NO_WARNINGS", + "/MD", + ] + return _common_opts + select({ + "//conditions:default": _posix_opts, + "@mediapipe//mediapipe:windows": _msvc_opts, + }) + +def halide_language_linkopts(): + _linux_opts = [ + "-ldl", + "-lpthread", + "-lz", + "-rdynamic", + ] + _osx_opts = [ + "-lz", + "-Wl,-stack_size", + "-Wl,1000000", + ] + _msvc_opts = [] + return select({ + "//conditions:default": _linux_opts, + "@mediapipe//mediapipe:macos": _osx_opts, + "@mediapipe//mediapipe:windows": _msvc_opts, + }) + +def halide_runtime_linkopts(): + """ Return the linkopts needed when linking against halide_library_runtime. + + Returns: + List to be used for linkopts. + """ + _posix_opts = [ + "-ldl", + "-lpthread", + ] + _android_opts = [ + "-llog", + ] + _msvc_opts = [] + + return select({ + "//conditions:default": _posix_opts, + "@mediapipe//mediapipe:android": _android_opts, + "@mediapipe//mediapipe:windows": _msvc_opts, + }) + +# Map of halide-target-base -> config_settings +_HALIDE_TARGET_CONFIG_SETTINGS_MAP = { + # Android + "arm-32-android": ["@mediapipe//mediapipe:android_arm"], + "arm-64-android": ["@mediapipe//mediapipe:android_arm64"], + "x86-32-android": ["@mediapipe//mediapipe:android_x86"], + "x86-64-android": ["@mediapipe//mediapipe:android_x86_64"], + # iOS + "arm-32-ios": ["@mediapipe//mediapipe:ios_armv7"], + "arm-64-ios": ["@mediapipe//mediapipe:ios_arm64", "@mediapipe//mediapipe:ios_arm64e"], + # OSX (or iOS simulator) + "x86-64-osx": [ + "@mediapipe//mediapipe:macos_x86_64", + "@mediapipe//mediapipe:ios_x86_64", + # TODO: these should map to "x86-32-osx", but that causes Kokoro to fail. + "@mediapipe//mediapipe:macos_i386", + "@mediapipe//mediapipe:ios_i386", + ], + "arm-64-osx": ["@mediapipe//mediapipe:macos_arm64"], + # Windows + "x86-64-windows": ["@mediapipe//mediapipe:windows"], + # Linux + "x86-64-linux": ["//conditions:default"], +} + +_HALIDE_TARGET_MAP_DEFAULT = { + "x86-64-linux": [ + "x86-64-linux-sse41-avx-avx2-fma", + "x86-64-linux-sse41", + "x86-64-linux", + ], + "x86-64-osx": [ + "x86-64-osx-sse41-avx-avx2-fma", + "x86-64-osx-sse41", + "x86-64-osx", + ], + "x86-64-windows": [ + "x86-64-windows-sse41-avx-avx2-fma", + "x86-64-windows-sse41", + "x86-64-windows", + ], +} + +def halide_library_default_target_map(): + return _HALIDE_TARGET_MAP_DEFAULT + +# Alphabetizes the features part of the target to make sure they always match no +# matter the concatenation order of the target string pieces. +def _canonicalize_target(halide_target): + if halide_target == "host": + return halide_target + if "," in halide_target: + fail("Multitarget may not be specified here") + tokens = halide_target.split("-") + if len(tokens) < 3: + fail("Illegal target: %s" % halide_target) + + # rejoin the tokens with the features sorted + return "-".join(tokens[0:3] + sorted(tokens[3:])) + +# Converts comma and dash separators to underscore and alphabetizes +# the features part of the target to make sure they always match no +# matter the concatenation order of the target string pieces. +def _halide_target_to_bazel_rule_name(multitarget): + subtargets = multitarget.split(",") + subtargets = [_canonicalize_target(st).replace("-", "_") for st in subtargets] + return "_".join(subtargets) + +# The second argument is True if there is a separate file generated +# for each subtarget of a multitarget output, False if not. The third +# argument is True if the output is a directory (vs. a single file). +# The fourth argument is a list of output group(s) that the files should +# be added to. + +_is_multi = True +_is_single = False +_is_file = False + +_output_extensions = { + "assembly": ("s", _is_multi, _is_file, []), + "bitcode": ("bc", _is_multi, _is_file, ["generated_bitcode"]), + "c_header": ("h", _is_single, _is_file, ["generated_headers"]), + "c_source": ("halide_generated.cpp", _is_multi, _is_file, []), + "compiler_log": ("halide_compiler_log", _is_single, _is_file, ["generated_object", "generated_compiler_log"]), + "cpp_stub": ("stub.h", _is_single, _is_file, []), + "featurization": ("featurization", _is_multi, _is_file, []), + "llvm_assembly": ("ll", _is_multi, _is_file, []), + "object": ("o", _is_single, _is_file, ["generated_object"]), + "python_extension": ("py.cpp", _is_single, _is_file, []), + "registration": ("registration.cpp", _is_single, _is_file, ["generated_registration"]), + "schedule": ("schedule.h", _is_single, _is_file, []), + "static_library": ("a", _is_single, _is_file, ["generated_object"]), + "stmt": ("stmt", _is_multi, _is_file, []), + "stmt_html": ("stmt.html", _is_multi, _is_file, []), +} + +def _add_output_file(f, fmt, output_files, output_dict, verbose_extra_outputs, verbose_output_paths): + if fmt in verbose_extra_outputs: + verbose_output_paths.append(f.path) + output_files.append(f) + if fmt in _output_extensions: + for group in _output_extensions[fmt][3]: + output_dict.setdefault(group, []).append(f) + +HalideFunctionNameInfo = provider(fields = ["function_name"]) +HalideGeneratorBinaryInfo = provider(fields = ["generator_binary"]) +HalideGeneratorNameInfo = provider(fields = ["generator_name_"]) +HalideGeneratorParamsInfo = provider(fields = ["generator_params"]) +HalideLibraryNameInfo = provider(fields = ["library_name"]) +HalideTargetFeaturesInfo = provider(fields = ["target_features"]) + +def _gengen_closure_impl(ctx): + return [ + HalideGeneratorBinaryInfo(generator_binary = ctx.attr.generator_binary), + HalideGeneratorNameInfo(generator_name_ = ctx.attr.generator_name_), + ] + +_gengen_closure = rule( + implementation = _gengen_closure_impl, + attrs = { + "generator_binary": attr.label( + executable = True, + allow_files = True, + mandatory = True, + cfg = "exec", + ), + # "generator_name" is apparently reserved by Bazel for attrs in rules + "generator_name_": attr.string(mandatory = True), + }, + provides = [HalideGeneratorBinaryInfo, HalideGeneratorNameInfo], +) + +def _halide_library_instance_impl(ctx): + generator_binary = ctx.attr.generator_closure[HalideGeneratorBinaryInfo].generator_binary if ctx.attr.generator_closure else "" + generator_name = ctx.attr.generator_closure[HalideGeneratorNameInfo].generator_name_ if ctx.attr.generator_closure else "" + return [ + HalideFunctionNameInfo(function_name = ctx.attr.function_name), + HalideGeneratorBinaryInfo(generator_binary = generator_binary), + HalideGeneratorNameInfo(generator_name_ = generator_name), + HalideGeneratorParamsInfo(generator_params = ctx.attr.generator_params), + HalideLibraryNameInfo(library_name = ctx.attr.library_name), + HalideTargetFeaturesInfo(target_features = ctx.attr.target_features), + ] + +_halide_library_instance = rule( + implementation = _halide_library_instance_impl, + attrs = { + "function_name": attr.string(), + "generator_closure": attr.label( + cfg = "exec", + providers = [HalideGeneratorBinaryInfo, HalideGeneratorNameInfo], + ), + "generator_params": attr.string_list(), + "library_name": attr.string(), + "target_features": attr.string_list(), + }, + provides = [ + HalideFunctionNameInfo, + HalideGeneratorBinaryInfo, + HalideGeneratorNameInfo, + HalideGeneratorParamsInfo, + HalideLibraryNameInfo, + HalideTargetFeaturesInfo, + ], +) + +def _gengen_impl(ctx): + if _has_dupes(ctx.attr.requested_outputs): + fail("Duplicate values in outputs: " + str(ctx.attr.requested_outputs)) + + function_name = ctx.attr.function_name[HalideFunctionNameInfo].function_name if ctx.attr.function_name else "" + generator_binary = ctx.attr.generator_binary[HalideGeneratorBinaryInfo].generator_binary if ctx.attr.generator_binary else "" + generator_name_ = ctx.attr.generator_name_[HalideGeneratorNameInfo].generator_name_ if ctx.attr.generator_name_ else "" + generator_params = ctx.attr.generator_params[HalideGeneratorParamsInfo].generator_params if ctx.attr.generator_params else [] + library_name = ctx.attr.library_name[HalideLibraryNameInfo].library_name if ctx.attr.library_name else "" + target_features = ctx.attr.target_features[HalideTargetFeaturesInfo].target_features if ctx.attr.target_features else [] + + for gp in generator_params: + if " " in gp: + fail("%s: Entries in generator_params must not contain spaces." % library_name) + + # Escape backslashes and double quotes. + generator_params = [gp.replace("\\", '\\\\"').replace('"', '\\"') for gp in generator_params] + + execution_requirements = {} + + # --- Calculate the output type(s) we're going to produce (and which ones should be verbose) + quiet_extra_outputs = [] + verbose_extra_outputs = [] + if ctx.attr.consider_halide_extra_outputs: + if "halide_extra_outputs" in ctx.var: + verbose_extra_outputs = ctx.var.get("halide_extra_outputs", "").split(",") + if "halide_extra_outputs_quiet" in ctx.var: + quiet_extra_outputs = ctx.var.get("halide_extra_outputs_quiet", "").split(",") + requested_outputs = sorted(collections.uniq(ctx.attr.requested_outputs + + verbose_extra_outputs + + quiet_extra_outputs)) + + # --- Assemble halide_target, adding extra features if necessary + base_target = ctx.attr.halide_base_target + if "," in base_target: + fail("halide_base_target should never be a multitarget") + if len(base_target.split("-")) != 3: + fail("halide_base_target should have exactly 3 components") + + target_features = target_features + ctx.var.get("halide_target_features", "").split(",") + + if "no_runtime" in target_features: + fail("Specifying 'no_runtime' in halide_target_features is not supported; " + + "please add 'add_halide_runtime_deps = False' to the halide_library() rule instead.") + + for san in ["asan", "msan", "tsan"]: + if san in target_features: + fail("halide_library doesn't support '%s' in halide_target_features; please build with --config=%s instead." % (san, san)) + + # Append the features common to everything. + target_features.append("c_plus_plus_name_mangling") + target_features.append("no_runtime") + + # Make it all neat and tidy. + target_features = sorted(collections.uniq(target_features)) + + # Get the multitarget list (if any) from halide_target_map + halide_targets = ctx.attr.halide_target_map.get(base_target, [base_target]) + + # Add the extra features to all of them + halide_targets = _add_features_to_all(halide_targets, target_features) + + leaf_name = ctx.attr.filename.split("/")[-1] + + output_files = [] + output_dict = {} + verbose_output_paths = [] + inputs = [] + + env = { + "HL_DEBUG_CODEGEN": str(ctx.var.get("halide_debug_codegen", 0)), + # --define halide_llvm_args=-time-passes is a typical usage + "HL_LLVM_ARGS": str(ctx.var.get("halide_llvm_args", "")), + } + + be_very_quiet = ctx.var.get("halide_experimental_quiet", False) # I'm hunting wabbit... + + # --- Calculate the final set of output files + for fmt in requested_outputs: + if fmt not in _output_extensions: + fail("Unknown Halide output '%s'; known outputs are %s" % + (fmt, sorted(_output_extensions.keys()))) + ext, is_multiple, is_dir, _ = _output_extensions[fmt] + + # Special-case Windows file extensions + if "windows" in halide_targets[-1]: + if ext == "o": + ext = "obj" + if ext == "a": + ext = "lib" + if is_multiple and len(halide_targets) > 1: + for h in halide_targets: + suffix = _canonicalize_target(h) + name = "%s-%s.%s" % (ctx.attr.filename, suffix, ext) + f = ctx.actions.declare_directory(name) if is_dir else ctx.actions.declare_file(name) + _add_output_file(f, fmt, output_files, output_dict, verbose_extra_outputs, verbose_output_paths) + else: + name = "%s.%s" % (ctx.attr.filename, ext) + f = ctx.actions.declare_directory(name) if is_dir else ctx.actions.declare_file(name) + _add_output_file(f, fmt, output_files, output_dict, verbose_extra_outputs, verbose_output_paths) + + # --- Progress message(s), including log info about any 'extra' files being output due to --define halide_extra_output + progress_message = "Executing generator %s with target (%s) args (%s)." % ( + generator_name_, + ",".join(halide_targets), + " ".join(generator_params), + ) + + for f in output_files: + if any([f.path.endswith(suf) for suf in [".h", ".a", ".o", ".lib", ".registration.cpp", ".bc", ".halide_compiler_log"]]): + continue + + # If an extra output was specified via --define halide_extra_outputs=foo on the command line, + # add to the progress message (so that it is ephemeral and doesn't clog stdout). + # + # (Trailing space is intentional since Starlark will append a period to the end, + # making copy-n-paste harder than it might otherwise be...) + if not be_very_quiet: + extra_msg = "Emitting extra Halide output: %s " % f.path + progress_message += "\n" + extra_msg + if f.path in verbose_output_paths: + # buildifier: disable=print + print(extra_msg) + + # --- Construct the arguments list for the Generator + arguments = ctx.actions.args() + arguments.add("-o", output_files[0].dirname) + if ctx.attr.generate_runtime: + arguments.add("-r", leaf_name) + if len(halide_targets) > 1: + fail("Only one halide_target allowed when using generate_runtime") + if function_name: + fail("halide_function_name not allowed when using generate_runtime") + else: + arguments.add("-g", generator_name_) + arguments.add("-n", leaf_name) + if function_name: + arguments.add("-f", function_name) + + if requested_outputs: + arguments.add_joined("-e", requested_outputs, join_with = ",") + + # Can't use add_joined(), as it will insert a space after target= + arguments.add("target=%s" % (",".join(halide_targets))) + if generator_params: + for p in generator_params: + for s in ["target"]: + if p.startswith("%s=" % s): + fail("You cannot specify %s in the generator_params parameter in bazel." % s) + arguments.add_all(generator_params) + + show_gen_arg = ctx.var.get("halide_show_generator_command", "") + + # If it's an exact match of a fully qualified path, show just that one. + # If it's * or "all", match everything. + if library_name and show_gen_arg in [library_name, "all", "*"] and not ctx.attr.generate_runtime: + # The 'Args' object can be printed, but can't be usefully converted to a string, or iterated, + # so we'll reproduce the logic here. We'll also take the opportunity to add or augment + # some args to be more useful to whoever runs it (eg, add `-v=1`, add some output files). + sg_args = ["-v", "1"] + sg_args += ["-o", "/tmp"] + sg_args += ["-g", generator_name_] + sg_args += ["-n", leaf_name] + if function_name: + sg_args += ["-f", function_name] + if requested_outputs: + # Ensure that several commonly-useful output are added + ro = sorted(collections.uniq(requested_outputs + ["stmt", "assembly", "llvm_assembly"])) + sg_args += ["-e", ",".join(ro)] + sg_args.append("target=%s" % (",".join(halide_targets))) + + if generator_params: + sg_args += generator_params + + # buildifier: disable=print + print( + "\n\nTo locally run the Generator for", + library_name, + "use the command:\n\n", + "bazel run -c opt", + generator_binary.label, + "--", + " ".join(sg_args), + "\n\n", + ) + + # Finally... run the Generator. + ctx.actions.run( + execution_requirements = execution_requirements, + arguments = [arguments], + env = env, + executable = generator_binary.files_to_run.executable, + mnemonic = "ExecuteHalideGenerator", + inputs = depset(direct = inputs), + outputs = output_files, + progress_message = progress_message, + exec_group = "generator", + ) + + return [ + DefaultInfo(files = depset(direct = output_files)), + OutputGroupInfo(**output_dict), + ] + +_gengen = rule( + implementation = _gengen_impl, + attrs = { + "consider_halide_extra_outputs": attr.bool(), + "filename": attr.string(), + "generate_runtime": attr.bool(default = False), + "generator_binary": attr.label( + cfg = "exec", + providers = [HalideGeneratorBinaryInfo], + ), + # "generator_name" is apparently reserved by Bazel for attrs in rules + "generator_name_": attr.label( + cfg = "exec", + providers = [HalideGeneratorNameInfo], + ), + "halide_base_target": attr.string(), + "function_name": attr.label( + cfg = "target", + providers = [ + HalideFunctionNameInfo, + ], + ), + "generator_params": attr.label( + cfg = "target", + providers = [ + HalideGeneratorParamsInfo, + ], + ), + "library_name": attr.label( + cfg = "target", + providers = [ + HalideLibraryNameInfo, + ], + ), + "target_features": attr.label( + cfg = "target", + providers = [ + HalideTargetFeaturesInfo, + ], + ), + "halide_target_map": attr.string_list_dict(), + "requested_outputs": attr.string_list(), + "_cc_toolchain": attr.label(default = "@bazel_tools//tools/cpp:current_cc_toolchain"), + }, + fragments = ["cpp"], + output_to_genfiles = True, + toolchains = use_cpp_toolchain(), + exec_groups = { + "generator": exec_group(), + }, +) + +def _add_target_features(target, features): + if "," in target: + fail("Cannot use multitarget here") + new_target = target.split("-") + for f in features: + if f and f not in new_target: + new_target.append(f) + return "-".join(new_target) + +def _add_features_to_all(halide_targets, features): + return [_canonicalize_target(_add_target_features(t, features)) for t in halide_targets] + +def _has_dupes(some_list): + clean = collections.uniq(some_list) + return sorted(some_list) != sorted(clean) + +# Target features which do not affect runtime compatibility. +_IRRELEVANT_FEATURES = collections.uniq([ + "arm_dot_prod", + "arm_fp16", + "c_plus_plus_name_mangling", + "check_unsafe_promises", + "embed_bitcode", + "enable_llvm_loop_opt", + "large_buffers", + "no_asserts", + "no_bounds_query", + "profile", + "strict_float", + "sve", + "sve2", + "trace_loads", + "trace_pipeline", + "trace_realizations", + "trace_stores", + "user_context", + "wasm_sat_float_to_int", + "wasm_signext", + "wasm_simd128", +]) + +def _discard_irrelevant_features(halide_target_features = []): + return sorted(collections.uniq([f for f in halide_target_features if f not in _IRRELEVANT_FEATURES])) + +def _halide_library_runtime_target_name(halide_target_features = []): + return "_".join(["halide_library_runtime"] + _discard_irrelevant_features(halide_target_features)) + +def _define_halide_library_runtime( + halide_target_features = [], + compatible_with = []): + target_name = _halide_library_runtime_target_name(halide_target_features) + + if not native.existing_rule("halide_library_runtime.generator"): + halide_generator( + name = "halide_library_runtime.generator", + srcs = [], + deps = [], + visibility = ["//visibility:private"], + ) + condition_deps = {} + for base_target, cfgs in _HALIDE_TARGET_CONFIG_SETTINGS_MAP.items(): + target_features = _discard_irrelevant_features(halide_target_features) + halide_target_name = _halide_target_to_bazel_rule_name(base_target) + gengen_name = "%s_%s" % (halide_target_name, target_name) + + _halide_library_instance( + name = "%s.library_instance" % gengen_name, + compatible_with = compatible_with, + function_name = "", + generator_closure = ":halide_library_runtime.generator_closure", + generator_params = [], + library_name = "", + target_features = target_features, + visibility = ["//visibility:private"], + ) + hl_instance = ":%s.library_instance" % gengen_name + + _gengen( + name = gengen_name, + compatible_with = compatible_with, + filename = "%s/%s" % (halide_target_name, target_name), + generate_runtime = True, + generator_binary = hl_instance, + generator_name_ = hl_instance, + halide_base_target = base_target, + requested_outputs = ["object"], + tags = ["manual"], + target_features = hl_instance, + visibility = ["@halide//:__subpackages__"], + ) + for cfg in cfgs: + condition_deps[cfg] = [":%s" % gengen_name] + + deps = [] + native.cc_library( + name = target_name, + compatible_with = compatible_with, + srcs = select(condition_deps), + linkopts = halide_runtime_linkopts(), + tags = ["manual"], + deps = deps, + visibility = ["//visibility:public"], + ) + + return target_name + +def _standard_library_runtime_features(): + _standard_features = [ + [], + ["cuda"], + ["metal"], + ["opencl"], + ["openglcompute"], + ["openglcompute", "egl"], + ] + return [f for f in _standard_features] + [f + ["debug"] for f in _standard_features] + +def _standard_library_runtime_names(): + return collections.uniq([_halide_library_runtime_target_name(f) for f in _standard_library_runtime_features()]) + +def halide_library_runtimes(compatible_with = []): + unused = [ + _define_halide_library_runtime(f, compatible_with = compatible_with) + for f in _standard_library_runtime_features() + ] + unused = unused # unused variable + +def halide_generator( + name, + srcs, + compatible_with = [], + copts = [], + deps = [], + generator_name = "", + includes = [], + tags = [], + testonly = False, + visibility = None): + if not name.endswith(".generator"): + fail("halide_generator rules must end in .generator") + + basename = name[:-10] # strip ".generator" suffix + if not generator_name: + generator_name = basename + + # Note: This target is public, but should not be needed by the vast + # majority of users. Unless you are writing a custom Bazel rule that + # involves Halide generation, you most probably won't need to depend on + # this rule. + native.cc_binary( + name = name, + copts = copts + halide_language_copts(), + linkopts = halide_language_linkopts(), + compatible_with = compatible_with, + srcs = srcs, + deps = [ + "@halide//:gengen", + "@halide//:language", + ] + deps, + tags = ["manual"] + tags, + testonly = testonly, + visibility = ["//visibility:public"], + ) + + _gengen_closure( + name = "%s_closure" % name, + generator_binary = name, + generator_name_ = generator_name, + compatible_with = compatible_with, + testonly = testonly, + visibility = ["//visibility:private"], + ) + +# This rule exists to allow us to select() on halide_target_features. +def _select_halide_library_runtime_impl(ctx): + f = ctx.attr.halide_target_features + + standard_runtimes = {t.label.name: t for t in ctx.attr._standard_runtimes} + + f = sorted(_discard_irrelevant_features(collections.uniq(f))) + runtime_name = _halide_library_runtime_target_name(f) + if runtime_name not in standard_runtimes: + fail(("There is no Halide runtime available for the feature set combination %s. " + + "Please use contact information from halide-lang.org to contact the Halide " + + "team to add the right combination.") % str(f)) + + return standard_runtimes[runtime_name][CcInfo] + +_select_halide_library_runtime = rule( + implementation = _select_halide_library_runtime_impl, + attrs = { + "halide_target_features": attr.string_list(), + "_standard_runtimes": attr.label_list( + default = ["@halide//:%s" % n for n in _standard_library_runtime_names()], + providers = [CcInfo], + ), + }, + provides = [CcInfo], +) + +def halide_library_from_generator( + name, + generator, + add_halide_runtime_deps = True, + compatible_with = [], + deps = [], + function_name = None, + generator_params = [], + halide_target_features = [], + halide_target_map = halide_library_default_target_map(), + includes = [], + namespace = None, + tags = [], + testonly = False, + visibility = None): + if not function_name: + function_name = name + + if namespace: + function_name = "%s::%s" % (namespace, function_name) + + generator_closure = "%s_closure" % generator + + _halide_library_instance( + name = "%s.library_instance" % name, + compatible_with = compatible_with, + function_name = function_name, + generator_closure = generator_closure, + generator_params = generator_params, + library_name = "//%s:%s" % (native.package_name(), name), + target_features = halide_target_features, + testonly = testonly, + visibility = ["//visibility:private"], + ) + hl_instance = ":%s.library_instance" % name + + condition_deps = {} + for base_target, cfgs in _HALIDE_TARGET_CONFIG_SETTINGS_MAP.items(): + base_target_name = _halide_target_to_bazel_rule_name(base_target) + gengen_name = "%s_%s" % (base_target_name, name) + _gengen( + name = gengen_name, + compatible_with = compatible_with, + consider_halide_extra_outputs = True, + filename = "%s/%s" % (base_target_name, name), + function_name = hl_instance, + generator_binary = generator_closure, + generator_name_ = generator_closure, + generator_params = hl_instance, + halide_base_target = base_target, + halide_target_map = halide_target_map, + library_name = hl_instance, + requested_outputs = ["static_library"], + tags = ["manual"] + tags, + target_features = hl_instance, + testonly = testonly, + ) + for cfg in cfgs: + condition_deps[cfg] = [":%s" % gengen_name] + + # Use a canonical target to build CC, regardless of config detected + cc_base_target = "x86-64-linux" + + for output, target_name in [ + ("c_header", "%s_h" % name), + ("c_source", "%s_cc" % name), + ]: + _gengen( + name = target_name, + compatible_with = compatible_with, + filename = name, + function_name = hl_instance, + generator_binary = generator_closure, + generator_name_ = generator_closure, + generator_params = hl_instance, + halide_base_target = cc_base_target, + library_name = hl_instance, + requested_outputs = [output], + tags = ["manual"] + tags, + target_features = hl_instance, + testonly = testonly, + ) + + _select_halide_library_runtime( + name = "%s.halide_library_runtime_deps" % name, + halide_target_features = halide_target_features, + compatible_with = compatible_with, + tags = tags, + visibility = ["//visibility:private"], + ) + + native.filegroup( + name = "%s_object" % name, + srcs = select(condition_deps), + output_group = "generated_object", + visibility = ["//visibility:private"], + compatible_with = compatible_with, + tags = tags, + testonly = testonly, + ) + + native.cc_library( + name = name, + srcs = ["%s_object" % name], + hdrs = [ + ":%s_h" % name, + ], + deps = deps + + ["@halide//:runtime"] + # for HalideRuntime.h, etc + ([":%s.halide_library_runtime_deps" % name] if add_halide_runtime_deps else []), # for the runtime implementation + defines = ["HALIDE_FUNCTION_ATTRS=HALIDE_MUST_USE_RESULT"], + compatible_with = compatible_with, + includes = includes, + tags = tags, + testonly = testonly, + visibility = visibility, + linkstatic = 1, + ) + + # Return the fully-qualified built target name. + return "//%s:%s" % (native.package_name(), name) + +def halide_library( + name, + srcs = [], + add_halide_runtime_deps = True, + copts = [], + compatible_with = [], + filter_deps = [], + function_name = None, + generator_params = [], + generator_deps = [], + generator_name = None, + halide_target_features = [], + halide_target_map = halide_library_default_target_map(), + includes = [], + namespace = None, + tags = [], + testonly = False, + visibility = None): + if not srcs and not generator_deps: + fail("halide_library needs at least one of srcs or generator_deps to provide a generator") + + halide_generator( + name = "%s.generator" % name, + srcs = srcs, + compatible_with = compatible_with, + generator_name = generator_name, + deps = generator_deps, + includes = includes, + copts = copts, + tags = tags, + testonly = testonly, + visibility = visibility, + ) + + return halide_library_from_generator( + name = name, + generator = ":%s.generator" % name, + add_halide_runtime_deps = add_halide_runtime_deps, + compatible_with = compatible_with, + deps = filter_deps, + function_name = function_name, + generator_params = generator_params, + halide_target_features = halide_target_features, + halide_target_map = halide_target_map, + includes = includes, + namespace = namespace, + tags = tags, + testonly = testonly, + visibility = visibility, + )