diff --git a/mediapipe/framework/formats/BUILD b/mediapipe/framework/formats/BUILD index dd311fb46..26525b5da 100644 --- a/mediapipe/framework/formats/BUILD +++ b/mediapipe/framework/formats/BUILD @@ -489,9 +489,12 @@ cc_test( cc_library( name = "frame_buffer", + srcs = ["frame_buffer.cc"], hdrs = ["frame_buffer.h"], deps = [ "//mediapipe/framework/port:integral_types", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", ], ) diff --git a/mediapipe/framework/formats/frame_buffer.cc b/mediapipe/framework/formats/frame_buffer.cc new file mode 100644 index 000000000..930a3651a --- /dev/null +++ b/mediapipe/framework/formats/frame_buffer.cc @@ -0,0 +1,176 @@ +/* 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. +==============================================================================*/ + +#include "mediapipe/framework/formats/frame_buffer.h" + +#include "absl/status/status.h" +#include "absl/status/statusor.h" + +namespace mediapipe { + +namespace { + +// Returns whether the input `format` is a supported YUV format. +bool IsSupportedYuvFormat(FrameBuffer::Format format) { + return format == FrameBuffer::Format::kNV21 || + format == FrameBuffer::Format::kNV12 || + format == FrameBuffer::Format::kYV12 || + format == FrameBuffer::Format::kYV21; +} + +// Returns supported 1-plane FrameBuffer in YuvData structure. +absl::StatusOr GetYuvDataFromOnePlaneFrameBuffer( + const FrameBuffer& source) { + if (!IsSupportedYuvFormat(source.format())) { + return absl::InvalidArgumentError( + "The source FrameBuffer format is not part of YUV420 family."); + } + + FrameBuffer::YuvData result; + const int y_buffer_size = + source.plane(0).stride().row_stride_bytes * source.dimension().height; + const int uv_buffer_size = + ((source.plane(0).stride().row_stride_bytes + 1) / 2) * + ((source.dimension().height + 1) / 2); + result.y_buffer = source.plane(0).buffer(); + result.y_row_stride = source.plane(0).stride().row_stride_bytes; + result.uv_row_stride = result.y_row_stride; + + if (source.format() == FrameBuffer::Format::kNV21) { + result.v_buffer = result.y_buffer + y_buffer_size; + result.u_buffer = result.v_buffer + 1; + result.uv_pixel_stride = 2; + // If y_row_stride equals to the frame width and is an odd value, + // uv_row_stride = y_row_stride + 1, otherwise uv_row_stride = y_row_stride. + if (result.y_row_stride == source.dimension().width && + result.y_row_stride % 2 == 1) { + result.uv_row_stride = (result.y_row_stride + 1) / 2 * 2; + } + } else if (source.format() == FrameBuffer::Format::kNV12) { + result.u_buffer = result.y_buffer + y_buffer_size; + result.v_buffer = result.u_buffer + 1; + result.uv_pixel_stride = 2; + // If y_row_stride equals to the frame width and is an odd value, + // uv_row_stride = y_row_stride + 1, otherwise uv_row_stride = y_row_stride. + if (result.y_row_stride == source.dimension().width && + result.y_row_stride % 2 == 1) { + result.uv_row_stride = (result.y_row_stride + 1) / 2 * 2; + } + } else if (source.format() == FrameBuffer::Format::kYV21) { + result.u_buffer = result.y_buffer + y_buffer_size; + result.v_buffer = result.u_buffer + uv_buffer_size; + result.uv_pixel_stride = 1; + result.uv_row_stride = (result.y_row_stride + 1) / 2; + } else if (source.format() == FrameBuffer::Format::kYV12) { + result.v_buffer = result.y_buffer + y_buffer_size; + result.u_buffer = result.v_buffer + uv_buffer_size; + result.uv_pixel_stride = 1; + result.uv_row_stride = (result.y_row_stride + 1) / 2; + } + return result; +} + +// Returns supported 2-plane FrameBuffer in YuvData structure. +absl::StatusOr GetYuvDataFromTwoPlaneFrameBuffer( + const FrameBuffer& source) { + if (source.format() != FrameBuffer::Format::kNV12 && + source.format() != FrameBuffer::Format::kNV21) { + return absl::InvalidArgumentError("Unsupported YUV planar format."); + } + + FrameBuffer::YuvData result; + // Y plane + result.y_buffer = source.plane(0).buffer(); + // All plane strides + result.y_row_stride = source.plane(0).stride().row_stride_bytes; + result.uv_row_stride = source.plane(1).stride().row_stride_bytes; + result.uv_pixel_stride = 2; + + if (source.format() == FrameBuffer::Format::kNV12) { + // Y and UV interleaved format + result.u_buffer = source.plane(1).buffer(); + result.v_buffer = result.u_buffer + 1; + } else { + // Y and VU interleaved format + result.v_buffer = source.plane(1).buffer(); + result.u_buffer = result.v_buffer + 1; + } + return result; +} + +// Returns supported 3-plane FrameBuffer in YuvData structure. Note that NV21 +// and NV12 are included in the supported Yuv formats. Technically, NV21 and +// NV12 should not be described by the 3-plane format. Historically, NV21 is +// used loosely such that it can also be used to describe YV21 format. For +// backwards compatibility, FrameBuffer supports NV21/NV12 with 3-plane format +// but such usage is discouraged +absl::StatusOr GetYuvDataFromThreePlaneFrameBuffer( + const FrameBuffer& source) { + if (!IsSupportedYuvFormat(source.format())) { + return absl::InvalidArgumentError( + "The source FrameBuffer format is not part of YUV420 family."); + } + + if (source.plane(1).stride().row_stride_bytes != + source.plane(2).stride().row_stride_bytes || + source.plane(1).stride().pixel_stride_bytes != + source.plane(2).stride().pixel_stride_bytes) { + return absl::InternalError("Unsupported YUV planar format."); + } + FrameBuffer::YuvData result; + if (source.format() == FrameBuffer::Format::kNV21 || + source.format() == FrameBuffer::Format::kYV12) { + // Y follow by VU order. The VU chroma planes can be interleaved or + // planar. + result.y_buffer = source.plane(0).buffer(); + result.v_buffer = source.plane(1).buffer(); + result.u_buffer = source.plane(2).buffer(); + result.y_row_stride = source.plane(0).stride().row_stride_bytes; + result.uv_row_stride = source.plane(1).stride().row_stride_bytes; + result.uv_pixel_stride = source.plane(1).stride().pixel_stride_bytes; + } else { + // Y follow by UV order. The UV chroma planes can be interleaved or + // planar. + result.y_buffer = source.plane(0).buffer(); + result.u_buffer = source.plane(1).buffer(); + result.v_buffer = source.plane(2).buffer(); + result.y_row_stride = source.plane(0).stride().row_stride_bytes; + result.uv_row_stride = source.plane(1).stride().row_stride_bytes; + result.uv_pixel_stride = source.plane(1).stride().pixel_stride_bytes; + } + return result; +} + +} // namespace + +absl::StatusOr FrameBuffer::GetYuvDataFromFrameBuffer( + const FrameBuffer& source) { + if (!IsSupportedYuvFormat(source.format())) { + return absl::InvalidArgumentError( + "The source FrameBuffer format is not part of YUV420 family."); + } + + if (source.plane_count() == 1) { + return GetYuvDataFromOnePlaneFrameBuffer(source); + } else if (source.plane_count() == 2) { + return GetYuvDataFromTwoPlaneFrameBuffer(source); + } else if (source.plane_count() == 3) { + return GetYuvDataFromThreePlaneFrameBuffer(source); + } + return absl::InvalidArgumentError( + "The source FrameBuffer must be consisted by 1, 2, or 3 planes"); +} + +} // namespace mediapipe diff --git a/mediapipe/framework/formats/frame_buffer.h b/mediapipe/framework/formats/frame_buffer.h index ccc699724..32ba41a2d 100644 --- a/mediapipe/framework/formats/frame_buffer.h +++ b/mediapipe/framework/formats/frame_buffer.h @@ -1,4 +1,4 @@ -/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. +/* 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. @@ -19,6 +19,7 @@ limitations under the License. #include #include "absl/log/check.h" +#include "absl/status/statusor.h" #include "mediapipe/framework/port/integral_types.h" namespace mediapipe { @@ -118,6 +119,20 @@ class FrameBuffer { int Size() const { return width * height; } }; + // YUV data structure. + struct YuvData { + const uint8* y_buffer; + const uint8* u_buffer; + const uint8* v_buffer; + // Y buffer row stride in bytes. + int y_row_stride; + // U/V buffer row stride in bytes. + int uv_row_stride; + // U/V pixel stride in bytes. This is the distance between two consecutive + // u/v pixel values in a row. + int uv_pixel_stride; + }; + // Builds a FrameBuffer object from a row-major backing buffer. // // The FrameBuffer does not take ownership of the backing buffer. The caller @@ -150,6 +165,12 @@ class FrameBuffer { // Returns FrameBuffer format. Format format() const { return format_; } + // Returns YuvData which contains the Y, U, and V buffer and their + // stride info from the input `source` FrameBuffer which is in the YUV family + // formats (e.g NV12, NV21, YV12, and YV21). + static absl::StatusOr GetYuvDataFromFrameBuffer( + const FrameBuffer& source); + private: std::vector planes_; Dimension dimension_;