Add FrameBuffer format.

PiperOrigin-RevId: 504838580

mediapipe/framework/formats/BUILD

@@ -486,3 +486,15 @@ cc_test(
         "//mediapipe/gpu:disable_gpu": [],
     }),
 )
+
+cc_library(
+    name = "frame_buffer",
+    srcs = ["frame_buffer.cc"],
+    hdrs = ["frame_buffer.h"],
+    deps = [
+        "//mediapipe/framework/port:integral_types",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+    ],
+)

mediapipe/framework/formats/frame_buffer.cc

@@ -0,0 +1,176 @@
+/* Copyright 2022 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

@@ -0,0 +1,246 @@
+/* Copyright 2022 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.
+==============================================================================*/
+
+#ifndef MEDIAPIPE_FRAMEWORK_FORMATS_FRAME_BUFFER_H_
+#define MEDIAPIPE_FRAMEWORK_FORMATS_FRAME_BUFFER_H_
+
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/status/statusor.h"
+#include "mediapipe/framework/port/integral_types.h"
+
+namespace mediapipe {
+
+// A `FrameBuffer` provides a view into the provided backing buffer (e.g. camera
+// frame or still image) with buffer format information. FrameBuffer doesn't
+// take ownership of the provided backing buffer. The caller is responsible to
+// manage the backing buffer lifecycle for the lifetime of the FrameBuffer.
+//
+// Examples:
+//
+// // Create an metadata instance with no backing buffer.
+// auto buffer = FrameBuffer::Create(/*planes=*/{}, dimension, kRGBA,
+//                                   KTopLeft);
+//
+// // Create an RGBA instance with backing buffer on single plane.
+// FrameBuffer::Plane plane =
+//     {rgba_buffer, /*stride=*/{dimension.width * 4, 4}};
+// auto buffer = FrameBuffer::Create({plane}, dimension, kRGBA, kTopLeft);
+//
+// // Create an YUV instance with planar backing buffer.
+// FrameBuffer::Plane y_plane = {y_buffer, /*stride=*/{dimension.width , 1}};
+// FrameBuffer::Plane uv_plane = {u_buffer, /*stride=*/{dimension.width, 2}};
+// auto buffer = FrameBuffer::Create({y_plane, uv_plane}, dimension, kNV21,
+//                                   kLeftTop);
+class FrameBuffer {
+ public:
+  // Colorspace formats.
+  enum class Format {
+    kRGBA,
+    kRGB,
+    kNV12,
+    kNV21,
+    kYV12,
+    kYV21,
+    kGRAY,
+    kUNKNOWN
+  };
+
+  // Stride information.
+  struct Stride {
+    // The row stride in bytes. This is the distance between the start pixels of
+    // two consecutive rows in the image.
+    int row_stride_bytes;
+    // This is the distance between two consecutive pixel values in a row of
+    // pixels in bytes. It may be larger than the size of a single pixel to
+    // account for interleaved image data or padded formats.
+    int pixel_stride_bytes;
+
+    bool operator==(const Stride& other) const {
+      return row_stride_bytes == other.row_stride_bytes &&
+             pixel_stride_bytes == other.pixel_stride_bytes;
+    }
+
+    bool operator!=(const Stride& other) const { return !operator==(other); }
+  };
+
+  // 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;
+  };
+
+  // FrameBuffer content orientation follows EXIF specification. The name of
+  // each enum value defines the position of the 0th row and the 0th column of
+  // the image content. See http://jpegclub.org/exif_orientation.html for
+  // details.
+  enum class Orientation {
+    kTopLeft = 1,
+    kTopRight = 2,
+    kBottomRight = 3,
+    kBottomLeft = 4,
+    kLeftTop = 5,
+    kRightTop = 6,
+    kRightBottom = 7,
+    kLeftBottom = 8
+  };
+
+  // Plane encapsulates buffer and stride information.
+  struct Plane {
+    const uint8* buffer;
+    Stride stride;
+  };
+
+  // Dimension information for the whole frame or a cropped portion of it.
+  struct Dimension {
+    // The width dimension in pixel unit.
+    int width;
+    // The height dimension in pixel unit.
+    int height;
+
+    bool operator==(const Dimension& other) const {
+      return width == other.width && height == other.height;
+    }
+
+    bool operator!=(const Dimension& other) const {
+      return width != other.width || height != other.height;
+    }
+
+    bool operator>=(const Dimension& other) const {
+      return width >= other.width && height >= other.height;
+    }
+
+    bool operator<=(const Dimension& other) const {
+      return width <= other.width && height <= other.height;
+    }
+
+    // Swaps width and height.
+    void Swap() {
+      using std::swap;
+      swap(width, height);
+    }
+
+    // Returns area represented by width * height.
+    int Size() const { return width * height; }
+  };
+
+  // Factory method for creating a FrameBuffer object from row-major backing
+  // buffers.
+  static std::unique_ptr<FrameBuffer> Create(const std::vector<Plane>& planes,
+                                             Dimension dimension, Format format,
+                                             Orientation orientation) {
+    return absl::make_unique<FrameBuffer>(planes, dimension, format,
+                                          orientation);
+  }
+
+  // Factory method for creating a FrameBuffer object from row-major movable
+  // backing buffers.
+  static std::unique_ptr<FrameBuffer> Create(std::vector<Plane>&& planes,
+                                             Dimension dimension, Format format,
+                                             Orientation orientation) {
+    return absl::make_unique<FrameBuffer>(std::move(planes), dimension, format,
+                                          orientation);
+  }
+
+  // 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);
+
+  // Builds a FrameBuffer object from a row-major backing buffer.
+  //
+  // The FrameBuffer does not take ownership of the backing buffer. The backing
+  // buffer is read-only and the caller is responsible for maintaining the
+  // backing buffer lifecycle for the lifetime of FrameBuffer.
+  FrameBuffer(const std::vector<Plane>& planes, Dimension dimension,
+              Format format, Orientation orientation)
+      : planes_(planes),
+        dimension_(dimension),
+        format_(format),
+        orientation_(orientation) {}
+
+  // Builds a FrameBuffer object from a movable row-major backing buffer.
+  //
+  // The FrameBuffer does not take ownership of the backing buffer. The backing
+  // buffer is read-only and the caller is responsible for maintaining the
+  // backing buffer lifecycle for the lifetime of FrameBuffer.
+  FrameBuffer(std::vector<Plane>&& planes, Dimension dimension, Format format,
+              Orientation orientation)
+      : planes_(std::move(planes)),
+        dimension_(dimension),
+        format_(format),
+        orientation_(orientation) {}
+
+  // Copy constructor.
+  //
+  // FrameBuffer does not take ownership of the backing buffer. The copy
+  // constructor behaves the same way to only copy the buffer pointer and not
+  // take ownership of the backing buffer.
+  FrameBuffer(const FrameBuffer& frame_buffer) {
+    planes_.clear();
+    for (int i = 0; i < frame_buffer.plane_count(); i++) {
+      planes_.push_back(
+          FrameBuffer::Plane{.buffer = frame_buffer.plane(i).buffer,
+                             .stride = frame_buffer.plane(i).stride});
+    }
+    dimension_ = frame_buffer.dimension();
+    format_ = frame_buffer.format();
+    orientation_ = frame_buffer.orientation();
+  }
+
+  // Returns number of planes.
+  int plane_count() const { return planes_.size(); }
+
+  // Returns plane indexed by the input `index`.
+  Plane plane(int index) const {
+    if (index > -1 && static_cast<size_t>(index) < planes_.size()) {
+      return planes_[index];
+    }
+    return {};
+  }
+
+  // Returns FrameBuffer dimension.
+  Dimension dimension() const { return dimension_; }
+
+  // Returns FrameBuffer format.
+  Format format() const { return format_; }
+
+  // Returns FrameBuffer orientation.
+  Orientation orientation() const { return orientation_; }
+
+ private:
+  std::vector<Plane> planes_;
+  Dimension dimension_;
+  Format format_;
+  Orientation orientation_;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_FORMATS_FRAME_BUFFER_H_