From 4c874fe4cd7f8c2fe4afdf1ac7630450264c3eba Mon Sep 17 00:00:00 2001
From: Camillo Lugaresi <camillol@google.com>
Date: Tue, 15 Nov 2022 18:31:27 -0800
Subject: [PATCH] Allow conversion of GlTextureBuffer to CVPixelBufferRef

This means that, if an iOS application sends in a GlTextureBuffer but expects a CVPixelBufferRef as output, everything will work even if the graph just forwards the same input. Also, access by Metal calculators will also work transparently.

PiperOrigin-RevId: 488812748
---
 mediapipe/gpu/BUILD                | 27 ++++++++++++++++++++++++++-
 mediapipe/gpu/gl_texture_buffer.cc | 29 +++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD
index 9cb27d2f1..196de3076 100644
--- a/mediapipe/gpu/BUILD
+++ b/mediapipe/gpu/BUILD
@@ -226,7 +226,13 @@ cc_library(
         # depend on having an indirect image_frame dependency, need to be
         # fixed first.
         "//mediapipe/framework/formats:image_frame",
-    ],
+    ] + select({
+        "//conditions:default": [],
+        ":platform_ios_with_gpu": [
+            ":gl_texture_util",
+            ":gpu_buffer_storage_cv_pixel_buffer",
+        ],
+    }),
 )
 
 cc_library(
@@ -344,6 +350,25 @@ cc_library(
     ],
 )
 
+mediapipe_cc_test(
+    name = "gpu_buffer_storage_cv_pixel_buffer_test",
+    size = "small",
+    timeout = "moderate",
+    srcs = ["gpu_buffer_storage_cv_pixel_buffer_test.cc"],
+    platforms = ["ios"],
+    deps = [
+        ":gl_texture_buffer",
+        ":gl_texture_util",
+        ":gpu_buffer",
+        ":gpu_buffer_storage_cv_pixel_buffer",
+        ":gpu_test_base",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/tool:test_util",
+        "//mediapipe/objc:util",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
 cc_library(
     name = "cv_texture_cache_manager",
     srcs = ["cv_texture_cache_manager.cc"],
diff --git a/mediapipe/gpu/gl_texture_buffer.cc b/mediapipe/gpu/gl_texture_buffer.cc
index fbb91a8f5..4c2f15a8d 100644
--- a/mediapipe/gpu/gl_texture_buffer.cc
+++ b/mediapipe/gpu/gl_texture_buffer.cc
@@ -18,6 +18,11 @@
 #include "mediapipe/gpu/gl_texture_view.h"
 #include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
 
+#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+#include "mediapipe/gpu/gl_texture_util.h"
+#include "mediapipe/gpu/gpu_buffer_storage_cv_pixel_buffer.h"
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+
 namespace mediapipe {
 
 std::unique_ptr<GlTextureBuffer> GlTextureBuffer::Wrap(
@@ -380,4 +385,28 @@ static auto kConverterRegistration2 =
         .RegisterConverter<GpuBufferStorageImageFrame, GlTextureBuffer>(
             ConvertFromImageFrame);
 
+#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+
+static std::shared_ptr<GpuBufferStorageCvPixelBuffer> ConvertToCvPixelBuffer(
+    std::shared_ptr<GlTextureBuffer> buf) {
+  auto output = absl::make_unique<GpuBufferStorageCvPixelBuffer>(
+      buf->width(), buf->height(), buf->format());
+  buf->GetProducerContext()->Run([buf, &output] {
+    TempGlFramebuffer framebuffer;
+    auto src = buf->GetReadView(internal::types<GlTextureView>{}, nullptr, 0);
+    auto dst =
+        output->GetWriteView(internal::types<GlTextureView>{}, nullptr, 0);
+    CopyGlTexture(src, dst);
+    glFlush();
+  });
+  return output;
+}
+
+static auto kConverterRegistrationCvpb =
+    internal::GpuBufferStorageRegistry::Get()
+        .RegisterConverter<GlTextureBuffer, GpuBufferStorageCvPixelBuffer>(
+            ConvertToCvPixelBuffer);
+
+#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+
 }  // namespace mediapipe