From 53fa35e40c00f2e58b6ef75d75fa6c0ee15e4c09 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 4 Apr 2023 01:16:50 -0700 Subject: [PATCH] Add FrameBuffer view on ImageFrame. PiperOrigin-RevId: 521689386 --- mediapipe/gpu/BUILD | 3 + mediapipe/gpu/frame_buffer_view.h | 37 +++ .../gpu/gpu_buffer_storage_image_frame.cc | 71 ++++++ .../gpu/gpu_buffer_storage_image_frame.h | 24 +- mediapipe/gpu/gpu_buffer_storage_yuv_image.cc | 228 ++++++++++++++++++ mediapipe/gpu/gpu_buffer_storage_yuv_image.h | 84 +++++++ 6 files changed, 446 insertions(+), 1 deletion(-) create mode 100644 mediapipe/gpu/frame_buffer_view.h create mode 100644 mediapipe/gpu/gpu_buffer_storage_image_frame.cc create mode 100644 mediapipe/gpu/gpu_buffer_storage_yuv_image.cc create mode 100644 mediapipe/gpu/gpu_buffer_storage_yuv_image.h diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD index ca2912ac3..c785e5624 100644 --- a/mediapipe/gpu/BUILD +++ b/mediapipe/gpu/BUILD @@ -423,12 +423,15 @@ cc_library( cc_library( name = "gpu_buffer_storage_image_frame", + srcs = ["gpu_buffer_storage_image_frame.cc"], hdrs = ["gpu_buffer_storage_image_frame.h"], visibility = ["//visibility:public"], deps = [ + ":frame_buffer_view", ":gpu_buffer_format", ":gpu_buffer_storage", ":image_frame_view", + "//mediapipe/framework/formats:frame_buffer", "//mediapipe/framework/formats:image_frame", ], ) diff --git a/mediapipe/gpu/frame_buffer_view.h b/mediapipe/gpu/frame_buffer_view.h new file mode 100644 index 000000000..76d773a5e --- /dev/null +++ b/mediapipe/gpu/frame_buffer_view.h @@ -0,0 +1,37 @@ +/* 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. +==============================================================================*/ + +#ifndef MEDIAPIPE_GPU_FRAME_BUFFER_VIEW_H_ +#define MEDIAPIPE_GPU_FRAME_BUFFER_VIEW_H_ + +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/gpu/gpu_buffer_storage.h" + +namespace mediapipe { +namespace internal { + +template <> +class ViewProvider { + public: + virtual ~ViewProvider() = default; + virtual std::shared_ptr GetReadView( + types) const = 0; + virtual std::shared_ptr GetWriteView(types) = 0; +}; + +} // namespace internal +} // namespace mediapipe + +#endif // MEDIAPIPE_GPU_FRAME_BUFFER_VIEW_H_ diff --git a/mediapipe/gpu/gpu_buffer_storage_image_frame.cc b/mediapipe/gpu/gpu_buffer_storage_image_frame.cc new file mode 100644 index 000000000..1cd661d37 --- /dev/null +++ b/mediapipe/gpu/gpu_buffer_storage_image_frame.cc @@ -0,0 +1,71 @@ +/* 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/gpu/gpu_buffer_storage_image_frame.h" + +#include +#include + +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/image_frame.h" + +namespace mediapipe { + +namespace { + +FrameBuffer::Format FrameBufferFormatForImageFrameFormat( + ImageFormat::Format format) { + switch (format) { + case ImageFormat::SRGB: + return FrameBuffer::Format::kRGB; + case ImageFormat::SRGBA: + return FrameBuffer::Format::kRGBA; + case ImageFormat::GRAY8: + return FrameBuffer::Format::kGRAY; + default: + return FrameBuffer::Format::kUNKNOWN; + } +} + +std::shared_ptr ImageFrameToFrameBuffer( + std::shared_ptr image_frame) { + FrameBuffer::Format format = + FrameBufferFormatForImageFrameFormat(image_frame->Format()); + CHECK(format != FrameBuffer::Format::kUNKNOWN) + << "Invalid format. Only SRGB, SRGBA and GRAY8 are supported."; + const FrameBuffer::Dimension dimension{/*width=*/image_frame->Width(), + /*height=*/image_frame->Height()}; + const FrameBuffer::Stride stride{ + /*row_stride_bytes=*/image_frame->WidthStep(), + /*pixel_stride_bytes=*/image_frame->ByteDepth() * + image_frame->NumberOfChannels()}; + const std::vector planes{ + {image_frame->MutablePixelData(), stride}}; + return std::make_shared(planes, dimension, format); +} + +} // namespace + +std::shared_ptr GpuBufferStorageImageFrame::GetReadView( + internal::types) const { + return ImageFrameToFrameBuffer(image_frame_); +} + +std::shared_ptr GpuBufferStorageImageFrame::GetWriteView( + internal::types) { + return ImageFrameToFrameBuffer(image_frame_); +} + +} // namespace mediapipe diff --git a/mediapipe/gpu/gpu_buffer_storage_image_frame.h b/mediapipe/gpu/gpu_buffer_storage_image_frame.h index ab547b9ea..542791f98 100644 --- a/mediapipe/gpu/gpu_buffer_storage_image_frame.h +++ b/mediapipe/gpu/gpu_buffer_storage_image_frame.h @@ -1,9 +1,26 @@ +/* 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. +==============================================================================*/ + #ifndef MEDIAPIPE_GPU_GPU_BUFFER_STORAGE_IMAGE_FRAME_H_ #define MEDIAPIPE_GPU_GPU_BUFFER_STORAGE_IMAGE_FRAME_H_ #include +#include "mediapipe/framework/formats/frame_buffer.h" #include "mediapipe/framework/formats/image_frame.h" +#include "mediapipe/gpu/frame_buffer_view.h" #include "mediapipe/gpu/gpu_buffer_format.h" #include "mediapipe/gpu/gpu_buffer_storage.h" #include "mediapipe/gpu/image_frame_view.h" @@ -13,7 +30,8 @@ namespace mediapipe { // Implements support for ImageFrame as a backing storage of GpuBuffer. class GpuBufferStorageImageFrame : public internal::GpuBufferStorageImpl< - GpuBufferStorageImageFrame, internal::ViewProvider> { + GpuBufferStorageImageFrame, internal::ViewProvider, + internal::ViewProvider> { public: explicit GpuBufferStorageImageFrame(std::shared_ptr image_frame) : image_frame_(image_frame) {} @@ -36,6 +54,10 @@ class GpuBufferStorageImageFrame internal::types) override { return image_frame_; } + std::shared_ptr GetReadView( + internal::types) const override; + std::shared_ptr GetWriteView( + internal::types) override; private: std::shared_ptr image_frame_; diff --git a/mediapipe/gpu/gpu_buffer_storage_yuv_image.cc b/mediapipe/gpu/gpu_buffer_storage_yuv_image.cc new file mode 100644 index 000000000..c7acd1340 --- /dev/null +++ b/mediapipe/gpu/gpu_buffer_storage_yuv_image.cc @@ -0,0 +1,228 @@ +/* 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/gpu/gpu_buffer_storage_yuv_image.h" + +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "libyuv/video_common.h" +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/image_frame.h" +#include "mediapipe/framework/formats/yuv_image.h" +#include "mediapipe/gpu/gpu_buffer_format.h" +#include "mediapipe/util/frame_buffer/frame_buffer_util.h" + +namespace mediapipe { + +namespace { + +// Default data alignment. +constexpr int kDefaultDataAligment = 16; + +GpuBufferFormat GpuBufferFormatForFourCC(libyuv::FourCC fourcc) { + switch (fourcc) { + case libyuv::FOURCC_NV12: + return GpuBufferFormat::kNV12; + case libyuv::FOURCC_NV21: + return GpuBufferFormat::kNV21; + case libyuv::FOURCC_YV12: + return GpuBufferFormat::kYV12; + case libyuv::FOURCC_I420: + return GpuBufferFormat::kI420; + default: + return GpuBufferFormat::kUnknown; + } +} + +libyuv::FourCC FourCCForGpuBufferFormat(GpuBufferFormat format) { + switch (format) { + case GpuBufferFormat::kNV12: + return libyuv::FOURCC_NV12; + case GpuBufferFormat::kNV21: + return libyuv::FOURCC_NV21; + case GpuBufferFormat::kYV12: + return libyuv::FOURCC_YV12; + case GpuBufferFormat::kI420: + return libyuv::FOURCC_I420; + default: + return libyuv::FOURCC_ANY; + } +} + +FrameBuffer::Format FrameBufferFormatForFourCC(libyuv::FourCC fourcc) { + switch (fourcc) { + case libyuv::FOURCC_NV12: + return FrameBuffer::Format::kNV12; + case libyuv::FOURCC_NV21: + return FrameBuffer::Format::kNV21; + case libyuv::FOURCC_YV12: + return FrameBuffer::Format::kYV12; + case libyuv::FOURCC_I420: + return FrameBuffer::Format::kYV21; + default: + return FrameBuffer::Format::kUNKNOWN; + } +} + +// Converts a YuvImage into a FrameBuffer that shares the same data buffers. +std::shared_ptr YuvImageToFrameBuffer( + std::shared_ptr yuv_image) { + FrameBuffer::Format format = FrameBufferFormatForFourCC(yuv_image->fourcc()); + FrameBuffer::Dimension dimension{/*width=*/yuv_image->width(), + /*height=*/yuv_image->height()}; + std::vector planes; + CHECK(yuv_image->mutable_data(0) != nullptr && yuv_image->stride(0) > 0) + << "Invalid YuvImage. Expected plane at index 0 to be non-null and have " + "stride > 0."; + planes.emplace_back( + yuv_image->mutable_data(0), + FrameBuffer::Stride{/*row_stride_bytes=*/yuv_image->stride(0), + /*pixel_stride_bytes=*/1}); + switch (format) { + case FrameBuffer::Format::kNV12: + case FrameBuffer::Format::kNV21: { + CHECK(yuv_image->mutable_data(1) != nullptr && yuv_image->stride(1) > 0) + << "Invalid YuvImage. Expected plane at index 1 to be non-null and " + "have stride > 0."; + planes.emplace_back( + yuv_image->mutable_data(1), + FrameBuffer::Stride{/*row_stride_bytes=*/yuv_image->stride(1), + /*pixel_stride_bytes=*/2}); + break; + } + case FrameBuffer::Format::kYV12: + case FrameBuffer::Format::kYV21: { + CHECK(yuv_image->mutable_data(1) != nullptr && yuv_image->stride(1) > 0 && + yuv_image->mutable_data(2) != nullptr && yuv_image->stride(2) > 0) + << "Invalid YuvImage. Expected planes at indices 1 and 2 to be " + "non-null and have stride > 0."; + planes.emplace_back( + yuv_image->mutable_data(1), + FrameBuffer::Stride{/*row_stride_bytes=*/yuv_image->stride(1), + /*pixel_stride_bytes=*/1}); + planes.emplace_back( + yuv_image->mutable_data(2), + FrameBuffer::Stride{/*row_stride_bytes=*/yuv_image->stride(2), + /*pixel_stride_bytes=*/1}); + break; + } + default: + LOG(FATAL) + << "Invalid format. Only FOURCC_NV12, FOURCC_NV21, FOURCC_YV12 and " + "FOURCC_I420 are supported."; + } + return std::make_shared(planes, dimension, format); +} + +// Converts a YUVImage into an ImageFrame with ImageFormat::SRGB format. +// Note that this requires YUV -> RGB conversion. +std::shared_ptr YuvImageToImageFrame( + std::shared_ptr yuv_image) { + auto yuv_buffer = YuvImageToFrameBuffer(yuv_image); + // Allocate the RGB ImageFrame to return. + auto image_frame = std::make_shared( + ImageFormat::SRGB, yuv_buffer->dimension().width, + yuv_buffer->dimension().height); + // Wrap it into a FrameBuffer + std::vector planes{ + {image_frame->MutablePixelData(), + {/*row_stride_bytes=*/image_frame->WidthStep(), + /*pixel_stride_bytes=*/image_frame->NumberOfChannels() * + image_frame->ChannelSize()}}}; + auto rgb_buffer = + FrameBuffer(planes, yuv_buffer->dimension(), FrameBuffer::Format::kRGB); + // Convert. + CHECK_OK(frame_buffer::Convert(*yuv_buffer, &rgb_buffer)); + return image_frame; +} + +} // namespace + +GpuBufferStorageYuvImage::GpuBufferStorageYuvImage( + std::shared_ptr yuv_image) { + CHECK(GpuBufferFormatForFourCC(yuv_image->fourcc()) != + GpuBufferFormat::kUnknown) + << "Invalid format. Only FOURCC_NV12, FOURCC_NV21, FOURCC_YV12 and " + "FOURCC_I420 are supported."; + yuv_image_ = yuv_image; +} + +GpuBufferStorageYuvImage::GpuBufferStorageYuvImage(int width, int height, + GpuBufferFormat format) { + libyuv::FourCC fourcc = FourCCForGpuBufferFormat(format); + int y_stride = std::ceil(1.0f * width / kDefaultDataAligment); + auto y_data = std::make_unique(y_stride * height); + switch (fourcc) { + case libyuv::FOURCC_NV12: + case libyuv::FOURCC_NV21: { + // Interleaved U/V planes, 2x2 downsampling. + int uv_width = 2 * std::ceil(0.5f * width); + int uv_height = std::ceil(0.5f * height); + int uv_stride = std::ceil(1.0f * uv_width / kDefaultDataAligment); + auto uv_data = std::make_unique(uv_stride * uv_height); + yuv_image_ = std::make_shared( + fourcc, std::move(y_data), y_stride, std::move(uv_data), uv_stride, + nullptr, 0, width, height); + break; + } + case libyuv::FOURCC_YV12: + case libyuv::FOURCC_I420: { + // Non-interleaved U/V planes, 2x2 downsampling. + int uv_width = std::ceil(0.5f * width); + int uv_height = std::ceil(0.5f * height); + int uv_stride = std::ceil(1.0f * uv_width / kDefaultDataAligment); + auto u_data = std::make_unique(uv_stride * uv_height); + auto v_data = std::make_unique(uv_stride * uv_height); + yuv_image_ = std::make_shared( + fourcc, std::move(y_data), y_stride, std::move(u_data), uv_stride, + std::move(v_data), uv_stride, width, height); + break; + } + default: + LOG(FATAL) + << "Invalid format. Only kNV12, kNV21, kYV12 and kYV21 are supported"; + } +} + +GpuBufferFormat GpuBufferStorageYuvImage::format() const { + return GpuBufferFormatForFourCC(yuv_image_->fourcc()); +} + +std::shared_ptr GpuBufferStorageYuvImage::GetReadView( + internal::types) const { + return YuvImageToFrameBuffer(yuv_image_); +} + +std::shared_ptr GpuBufferStorageYuvImage::GetWriteView( + internal::types) { + return YuvImageToFrameBuffer(yuv_image_); +} + +std::shared_ptr GpuBufferStorageYuvImage::GetReadView( + internal::types) const { + return YuvImageToImageFrame(yuv_image_); +} + +std::shared_ptr GpuBufferStorageYuvImage::GetWriteView( + internal::types) { + // Not supported on purpose: writes into the resulting ImageFrame cannot + // easily be ported back to the original YUV image. + LOG(FATAL) << "GetWriteView is not supported."; +} +} // namespace mediapipe diff --git a/mediapipe/gpu/gpu_buffer_storage_yuv_image.h b/mediapipe/gpu/gpu_buffer_storage_yuv_image.h new file mode 100644 index 000000000..6b34f4948 --- /dev/null +++ b/mediapipe/gpu/gpu_buffer_storage_yuv_image.h @@ -0,0 +1,84 @@ +/* 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 + +#include "mediapipe/framework/formats/frame_buffer.h" +#include "mediapipe/framework/formats/image_frame.h" +#include "mediapipe/framework/formats/yuv_image.h" +#include "mediapipe/gpu/frame_buffer_view.h" +#include "mediapipe/gpu/gpu_buffer_format.h" +#include "mediapipe/gpu/gpu_buffer_storage.h" +#include "mediapipe/gpu/image_frame_view.h" + +#ifndef MEDIAPIPE_GPU_GPU_BUFFER_STORAGE_YUV_IMAGE_H_ +#define MEDIAPIPE_GPU_GPU_BUFFER_STORAGE_YUV_IMAGE_H_ + +namespace mediapipe { + +namespace internal { + +template <> +class ViewProvider { + public: + virtual ~ViewProvider() = default; + virtual std::shared_ptr GetReadView( + types) const = 0; + virtual std::shared_ptr GetWriteView(types) = 0; +}; + +} // namespace internal + +// TODO: add support for I444. +class GpuBufferStorageYuvImage + : public internal::GpuBufferStorageImpl< + GpuBufferStorageYuvImage, internal::ViewProvider, + internal::ViewProvider, + internal::ViewProvider> { + public: + // Constructor from an existing YUVImage with FOURCC_NV12, FOURCC_NV21, + // FOURCC_YV12 or FOURCC_I420 format. + explicit GpuBufferStorageYuvImage(std::shared_ptr yuv_image); + // Constructor. Supported formats are kNV12, kNV21, kYV12 and kI420. + // Stride is set by default so that row boundaries align to 16 bytes. + GpuBufferStorageYuvImage(int width, int height, GpuBufferFormat format); + + int width() const override { return yuv_image_->width(); } + int height() const override { return yuv_image_->height(); } + GpuBufferFormat format() const override; + + std::shared_ptr GetReadView( + internal::types) const override { + return yuv_image_; + } + std::shared_ptr GetWriteView(internal::types) override { + return yuv_image_; + } + + std::shared_ptr GetReadView( + internal::types) const override; + std::shared_ptr GetWriteView( + internal::types) override; + std::shared_ptr GetReadView( + internal::types) const override; + std::shared_ptr GetWriteView( + internal::types) override; + + private: + std::shared_ptr yuv_image_; +}; +} // namespace mediapipe + +#endif // MEDIAPIPE_GPU_GPU_BUFFER_STORAGE_YUV_IMAGE_H_