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

Add build system for Halide and expose FrameBufferUtils.

Browse Source 
PiperOrigin-RevId: 515304264
This commit is contained in:
MediaPipe Team 2023-03-09 05:04:18 -08:00 committed by Copybara-Service
parent 5398b8881d
commit 39e2c8351f
36 changed files with 6501 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
WORKSPACE
View File

@ -543,3 +543,35 @@ external_files()
load("@//third_party:wasm_files.bzl", "wasm_files") load("@//third_party:wasm_files.bzl", "wasm_files")
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",
)

106
mediapipe/util/frame_buffer/BUILD Normal file
View File

@ -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",
],
)

40
mediapipe/util/frame_buffer/buffer_common.cc Normal file
View File

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

32
mediapipe/util/frame_buffer/buffer_common.h Normal file
View File

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

794
mediapipe/util/frame_buffer/frame_buffer_util.cc Normal file
View File

@ -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 <cstdint>
#include <memory>
#include <vector>
#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<YuvBuffer> CreateYuvBuffer(const FrameBuffer& buffer) {
ASSIGN_OR_RETURN(FrameBuffer::YuvData yuv_data,
FrameBuffer::GetYuvDataFromFrameBuffer(buffer));
return YuvBuffer(const_cast<uint8_t*>(yuv_data.y_buffer),
const_cast<uint8_t*>(yuv_data.u_buffer),
const_cast<uint8_t*>(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<GrayBuffer> CreateGrayBuffer(const FrameBuffer& buffer) {
if (buffer.plane_count() != 1) {
return absl::InternalError("Unsupported grayscale planar format.");
}
return GrayBuffer(const_cast<uint8_t*>(buffer.plane(0).buffer()),
buffer.dimension().width, buffer.dimension().height);
}
absl::StatusOr<RgbBuffer> 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<uint8_t*>(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<uint8_t*>(output_buffer->plane(0).buffer()));
} else {
// The y_buffer is padded. The conversion removes the padding.
uint8_t* gray_buffer =
const_cast<uint8_t*>(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<FrameBuffer> 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<FrameBuffer::Plane> planes{input_plane};
return std::make_shared<FrameBuffer>(planes, dimension,
FrameBuffer::Format::kRGBA);
}
std::shared_ptr<FrameBuffer> 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<FrameBuffer::Plane> planes{input_plane};
return std::make_shared<FrameBuffer>(planes, dimension,
FrameBuffer::Format::kRGB);
}
std::shared_ptr<FrameBuffer> 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<FrameBuffer::Plane> planes{input_plane};
return std::make_shared<FrameBuffer>(planes, dimension,
FrameBuffer::Format::kGRAY);
}
absl::StatusOr<std::shared_ptr<FrameBuffer>> 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<FrameBuffer::Plane> 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<FrameBuffer>(planes, dimension, format);
}
absl::StatusOr<std::shared_ptr<FrameBuffer>> 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<FrameBuffer::Plane> planes{plane};
return std::make_shared<FrameBuffer>(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<int> 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<const uint8_t*> 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<FrameBuffer::Dimension> 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

131
mediapipe/util/frame_buffer/frame_buffer_util.h Normal file
View File

@ -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 <cstdint>
#include <memory>
#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<FrameBuffer> CreateFromRgbaRawBuffer(
uint8_t* input, FrameBuffer::Dimension dimension,
FrameBuffer::Stride stride = kDefaultStride);
// Creates a FrameBuffer from raw RGB buffer and passing arguments.
std::shared_ptr<FrameBuffer> CreateFromRgbRawBuffer(
uint8_t* input, FrameBuffer::Dimension dimension,
FrameBuffer::Stride stride = kDefaultStride);
// Creates a FrameBuffer from raw grayscale buffer and passing arguments.
std::shared_ptr<FrameBuffer> CreateFromGrayRawBuffer(
uint8_t* input, FrameBuffer::Dimension dimension,
FrameBuffer::Stride stride = kDefaultStride);
// Creates a FrameBuffer from raw YUV buffer and passing arguments.
absl::StatusOr<std::shared_ptr<FrameBuffer>> 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<std::shared_ptr<FrameBuffer>> 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<int> GetPixelStrides(FrameBuffer::Format format);
// Returns the biplanar UV raw buffer for NV12/NV21 frame buffer.
absl::StatusOr<const uint8_t*> 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<FrameBuffer::Dimension> 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_

1176
mediapipe/util/frame_buffer/frame_buffer_util_test.cc Normal file
View File

File diff suppressed because it is too large Load Diff

95
mediapipe/util/frame_buffer/gray_buffer.cc Normal file
View File

@ -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 <utility>
#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<uint8_t>(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<float>(width()) / output->width(),
static_cast<float>(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

137
mediapipe/util/frame_buffer/gray_buffer.h Normal file
View File

@ -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 <memory>
#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<uint8_t[]> 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<uint8_t> buffer_;
};
} // namespace frame_buffer
} // namespace mediapipe
#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_GRAY_BUFFER_H_

223
mediapipe/util/frame_buffer/gray_buffer_test.cc Normal file
View File

@ -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 <utility>
#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

118
mediapipe/util/frame_buffer/halide/BUILD Normal file
View File

@ -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",
)

89
mediapipe/util/frame_buffer/halide/common.cc Normal file
View File

@ -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<int>((x + 0.5f) * fx),
Halide::cast<int>((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<int>(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<int>(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<int>(fx * 65536);
Halide::Expr xi = Halide::cast<int>(x * fx / 65536);
Halide::Expr xr = Halide::cast<uint16_t>(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<int>(fy * 65536);
Halide::Expr yi = Halide::cast<int>(y * fy / 65536);
Halide::Expr yr = Halide::cast<uint16_t>(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

61
mediapipe/util/frame_buffer/halide/common.h Normal file
View File

@ -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 <typename T>
Halide::Expr is_planar(const T& buffer) {
return buffer.dim(0).stride() == 1;
}
template <typename T>
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_

61
mediapipe/util/frame_buffer/halide/gray_flip_generator.cc Normal file
View File

@ -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<GrayFlip> {
public:
Var x{"x"}, y{"y"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
// Flip vertically if true; flips horizontally (mirroring) otherwise.
Input<bool> flip_vertical{"flip_vertical", false};
Output<Func> 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)

60
mediapipe/util/frame_buffer/halide/gray_resize_generator.cc Normal file
View File

@ -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<GrayResize> {
public:
Var x{"x"}, y{"y"};
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
Input<float> scale_x{"scale_x", 1.0f, 0.0f, 1024.0f};
Input<float> scale_y{"scale_y", 1.0f, 0.0f, 1024.0f};
Output<Func> 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<uint8_t>();
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)

63
mediapipe/util/frame_buffer/halide/gray_rotate_generator.cc Normal file
View File

@ -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<GrayRotate> {
public:
Var x{"x"}, y{"y"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
// Rotation angle in degrees counter-clockwise. Must be in {0, 90, 180, 270}.
Input<int> rotation_angle{"rotation_angle", 0};
Output<Func> 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)

84
mediapipe/util/frame_buffer/halide/rgb_flip_generator.cc Normal file
View File

@ -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<RgbFlip> {
public:
Var x{"x"}, y{"y"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 3>> src_rgb{"src_rgb"};
// Flip vertically if true; flips horizontally (mirroring) otherwise.
Input<bool> flip_vertical{"flip_vertical", false};
Output<Func> 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)

65
mediapipe/util/frame_buffer/halide/rgb_gray_generator.cc Normal file
View File

@ -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<RgbGray> {
public:
Var x{"x"}, y{"y"}, c{"c"};
Input<Buffer<uint8_t, 3>> src_rgb{"rgb"};
Output<Buffer<uint8_t, 2>> 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<int32_t>(r);
g = Halide::cast<int32_t>(g);
b = Halide::cast<int32_t>(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<uint8_t>(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)

85
mediapipe/util/frame_buffer/halide/rgb_resize_generator.cc Normal file
View File

@ -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<RgbResize> {
public:
Var x{"x"}, y{"y"};
Input<Buffer<uint8_t, 3>> src_rgb{"src_rgb"};
Input<float> scale_x{"scale_x", 1.0f, 0.0f, 1024.0f};
Input<float> scale_y{"scale_y", 1.0f, 0.0f, 1024.0f};
Output<Func> 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<uint8_t>();
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)

64
mediapipe/util/frame_buffer/halide/rgb_rgb_generator.cc Normal file
View File

@ -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 <cstdint>
#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<RgbRgb> {
public:
Var x{"x"}, y{"y"}, c{"c"};
Input<Buffer<uint8_t, 3>> src_rgb{"src_rgb"};
Output<Buffer<uint8_t, 3>> 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)

76
mediapipe/util/frame_buffer/halide/rgb_rotate_generator.cc Normal file
View File

@ -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<RgbRotate> {
public:
Var x{"x"}, y{"y"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 3>> src_rgb{"src_rgb"};
// Rotation angle in degrees counter-clockwise. Must be in {0, 90, 180, 270}.
Input<int> rotation_angle{"rotation_angle", 0};
Output<Func> 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)

101
mediapipe/util/frame_buffer/halide/rgb_yuv_generator.cc Normal file
View File

@ -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<RgbYuv> {
public:
Var x{"x"}, y{"y"}, c{"c"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 3>> src_rgb{"rgb"};
Output<Func> dst_y{"dst_y", UInt(8), 2};
Output<Func> 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<int32_t>(r);
g = Halide::cast<int32_t>(g);
b = Halide::cast<int32_t>(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<uint8_t>(yuv_tuple(x, y)[0]);
dst_uv(x, y, c) = Halide::saturating_cast<uint8_t>(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)

90
mediapipe/util/frame_buffer/halide/yuv_flip_generator.cc Normal file
View File

@ -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<YuvFlip> {
public:
Var x{"x"}, y{"y"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
Input<Buffer<uint8_t, 3>> src_uv{"src_uv"};
// Flip vertically if true; flips horizontally (mirroring) otherwise.
Input<bool> flip_vertical{"flip_vertical", false};
Output<Func> dst_y{"dst_y", UInt(8), 2};
Output<Func> 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)

91
mediapipe/util/frame_buffer/halide/yuv_resize_generator.cc Normal file
View File

@ -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<YuvResize> {
public:
Var x{"x"}, y{"y"};
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
Input<Buffer<uint8_t, 3>> src_uv{"src_uv"};
Input<float> scale_x{"scale_x", 1.0f, 0.0f, 1024.0f};
Input<float> scale_y{"scale_y", 1.0f, 0.0f, 1024.0f};
Output<Func> dst_y{"dst_y", UInt(8), 2};
Output<Func> 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<uint8_t>();
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)

110
mediapipe/util/frame_buffer/halide/yuv_rgb_generator.cc Normal file
View File

@ -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<YuvRgb> {
public:
Var x{"x"}, y{"y"}, c{"c"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
Input<Buffer<uint8_t, 3>> src_uv{"src_uv"};
Input<bool> halve{"halve", false};
Output<Func> 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<int32_t>(y);
u = Halide::cast<int32_t>(u) - 128;
v = Halide::cast<int32_t>(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<uint8_t>(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<uint8_t>();
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)

91
mediapipe/util/frame_buffer/halide/yuv_rotate_generator.cc Normal file
View File

@ -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<YuvRotate> {
public:
Var x{"x"}, y{"y"};
// Input<Buffer> because that allows us to apply constraints on stride, etc.
Input<Buffer<uint8_t, 2>> src_y{"src_y"};
Input<Buffer<uint8_t, 3>> src_uv{"src_uv"};
// Rotation angle in degrees counter-clockwise. Must be in {0, 90, 180, 270}.
Input<int> rotation_angle{"rotation_angle", 0};
Output<Func> dst_y{"dst_y", UInt(8), 2};
Output<Func> 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)

132
mediapipe/util/frame_buffer/rgb_buffer.cc Normal file
View File

@ -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 <utility>
#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<uint8_t>(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<float>(width()) / output->width(),
static_cast<float>(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<uint8_t>::make_interleaved(
data, width, height, channels);
}
} // namespace frame_buffer
} // namespace mediapipe

139
mediapipe/util/frame_buffer/rgb_buffer.h Normal file
View File

@ -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 <memory>
#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<uint8_t[]> owned_buffer_;
// Backing buffer: layout is always width x height x channel (interleaved).
Halide::Runtime::Buffer<uint8_t> buffer_;
};
} // namespace frame_buffer
} // namespace mediapipe
#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_RGB_BUFFER_H_

606
mediapipe/util/frame_buffer/rgb_buffer_test.cc Normal file
View File

@ -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 <utility>
#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<GrayBuffer&>(lhs).buffer()->host;
const uint8_t* converted_ptr = const_cast<GrayBuffer&>(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<RgbBuffer&>(lhs).buffer()->host;
const uint8_t* converted_ptr = const_cast<RgbBuffer&>(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<YuvBuffer&>(lhs).y_buffer()->host;
const uint8_t* converted_ptr = const_cast<YuvBuffer&>(rhs).y_buffer()->host;
if (!CompareArray(reference_ptr, converted_ptr, lhs.width(), lhs.height())) {
return false;
}
reference_ptr = const_cast<YuvBuffer&>(lhs).uv_buffer()->host;
converted_ptr = const_cast<YuvBuffer&>(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

152
mediapipe/util/frame_buffer/yuv_buffer.cc Normal file
View File

@ -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 <utility>
#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<uint8_t>(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<int32_t>(u_plane - v_plane)},
};
uv_buffer_ = Halide::Runtime::Buffer<uint8_t>(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<uint8_t>(data, width, height);
uint8_t* uv_data = data + (width * height);
switch (format) {
case NV21:
// Interleaved UV (actually VU order).
uv_buffer_ = Halide::Runtime::Buffer<uint8_t>::make_interleaved(
uv_data, (width + 1) / 2, (height + 1) / 2, 2);
break;
case YV12:
// Planar UV (actually VU order).
uv_buffer_ = Halide::Runtime::Buffer<uint8_t>(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<float>(width()) / output->width(),
static_cast<float>(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

160
mediapipe/util/frame_buffer/yuv_buffer.h Normal file
View File

@ -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 <memory>
#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<uint8_t[]> owned_buffer_;
// Y (luminance) backing buffer: layout is always width x height.
Halide::Runtime::Buffer<uint8_t> 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<uint8_t> uv_buffer_;
};
} // namespace frame_buffer
} // namespace mediapipe
#endif // MEDIAPIPE_UTIL_FRAME_BUFFER_YUV_BUFFER_H_

251
mediapipe/util/frame_buffer/yuv_buffer_test.cc Normal file
View File

@ -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 <utility>
#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

70
third_party/halide.BUILD vendored Normal file
View File

@ -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",
],
)

1
third_party/halide/BUILD vendored Normal file
View File

@ -0,0 +1 @@
# This empty BUILD file is required to make Bazel treat this directory as a package.

40
third_party/halide/BUILD.bazel vendored Normal file
View File

@ -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",
]
]

875
third_party/halide/halide.bzl vendored Normal file
View File

@ -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,
)