From b308c0dd5e114cbf803dd2864f67589be048b7a0 Mon Sep 17 00:00:00 2001
From: Camillo Lugaresi <camillol@google.com>
Date: Tue, 15 Nov 2022 17:08:37 -0800
Subject: [PATCH] Implement CVPixelBufferRef access as a view.

PiperOrigin-RevId: 488798216
---
 mediapipe/gpu/gpu_buffer.cc                   |  7 ++--
 mediapipe/gpu/gpu_buffer.h                    |  4 +++
 .../gpu/gpu_buffer_storage_cv_pixel_buffer.h  | 35 ++++++++++++++++++-
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/mediapipe/gpu/gpu_buffer.cc b/mediapipe/gpu/gpu_buffer.cc
index 35a73fd8..388960b1 100644
--- a/mediapipe/gpu/gpu_buffer.cc
+++ b/mediapipe/gpu/gpu_buffer.cc
@@ -93,8 +93,11 @@ internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
 #if !MEDIAPIPE_DISABLE_GPU
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer) {
-  auto p = buffer.internal_storage<GpuBufferStorageCvPixelBuffer>();
-  if (p) return **p;
+  if (buffer.GetStorageForView(
+          kTypeId<internal::ViewProvider<CVPixelBufferRef>>,
+          /*for_writing=*/false) != nullptr) {
+    return *buffer.GetReadView<CVPixelBufferRef>();
+  }
   return nullptr;
 }
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
diff --git a/mediapipe/gpu/gpu_buffer.h b/mediapipe/gpu/gpu_buffer.h
index ad5c130b..45146a32 100644
--- a/mediapipe/gpu/gpu_buffer.h
+++ b/mediapipe/gpu/gpu_buffer.h
@@ -179,6 +179,10 @@ class GpuBuffer {
   // This is mutable because view methods that do not change the contents may
   // still need to allocate new storages.
   mutable std::vector<std::shared_ptr<internal::GpuBufferStorage>> storages_;
+
+#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+  friend CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer);
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 };
 
 inline bool GpuBuffer::operator==(std::nullptr_t other) const {
diff --git a/mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.h b/mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.h
index 017771dc..e5bc5de4 100644
--- a/mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.h
+++ b/mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.h
@@ -12,10 +12,27 @@ namespace mediapipe {
 
 class GlContext;
 
+namespace internal {
+
+template <>
+class ViewProvider<CVPixelBufferRef> {
+ public:
+  virtual ~ViewProvider() = default;
+  virtual CFHolder<CVPixelBufferRef> GetReadView(
+      internal::types<CVPixelBufferRef>,
+      std::shared_ptr<GpuBuffer> gpu_buffer) const = 0;
+  virtual CFHolder<CVPixelBufferRef> GetWriteView(
+      internal::types<CVPixelBufferRef>,
+      std::shared_ptr<GpuBuffer> gpu_buffer) = 0;
+};
+
+}  // namespace internal
+
 class GpuBufferStorageCvPixelBuffer
     : public internal::GpuBufferStorageImpl<
           GpuBufferStorageCvPixelBuffer, internal::ViewProvider<GlTextureView>,
-          internal::ViewProvider<ImageFrame>>,
+          internal::ViewProvider<ImageFrame>,
+          internal::ViewProvider<CVPixelBufferRef>>,
       public CFHolder<CVPixelBufferRef> {
  public:
   using CFHolder<CVPixelBufferRef>::CFHolder;
@@ -44,6 +61,12 @@ class GpuBufferStorageCvPixelBuffer
   std::shared_ptr<ImageFrame> GetWriteView(
       internal::types<ImageFrame>,
       std::shared_ptr<GpuBuffer> gpu_buffer) override;
+  CFHolder<CVPixelBufferRef> GetReadView(
+      internal::types<CVPixelBufferRef>,
+      std::shared_ptr<GpuBuffer> gpu_buffer) const override;
+  CFHolder<CVPixelBufferRef> GetWriteView(
+      internal::types<CVPixelBufferRef>,
+      std::shared_ptr<GpuBuffer> gpu_buffer) override;
 
  private:
   GlTextureView GetTexture(std::shared_ptr<GpuBuffer> gpu_buffer, int plane,
@@ -51,6 +74,16 @@ class GpuBufferStorageCvPixelBuffer
   void ViewDoneWriting(const GlTextureView& view);
 };
 
+inline CFHolder<CVPixelBufferRef> GpuBufferStorageCvPixelBuffer::GetReadView(
+    internal::types<CVPixelBufferRef>,
+    std::shared_ptr<GpuBuffer> gpu_buffer) const {
+  return *this;
+}
+inline CFHolder<CVPixelBufferRef> GpuBufferStorageCvPixelBuffer::GetWriteView(
+    internal::types<CVPixelBufferRef>, std::shared_ptr<GpuBuffer> gpu_buffer) {
+  return *this;
+}
+
 namespace internal {
 // These functions enable backward-compatible construction of a GpuBuffer from
 // CVPixelBufferRef without having to expose that type in the main GpuBuffer