Internal change

PiperOrigin-RevId: 509338497

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

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<FrameBuffer::YuvData> 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<FrameBuffer::YuvData> 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<FrameBuffer::YuvData> 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::YuvData> 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

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 <vector>
 
 #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<YuvData> GetYuvDataFromFrameBuffer(
+      const FrameBuffer& source);
+
  private:
   std::vector<Plane> planes_;
   Dimension dimension_;