From 1df4511e9d5a8603835512318606262be949f9ff Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Fri, 27 Jan 2023 08:49:28 -0800
Subject: [PATCH] Add YuvImage as a GpuBuffer storage backend.

PiperOrigin-RevId: 505128789
---
 mediapipe/gpu/BUILD                | 29 +++++++++++++++++++++++++++++
 mediapipe/gpu/gpu_buffer_format.cc |  4 ++++
 mediapipe/gpu/gpu_buffer_format.h  | 12 ++++++++++++
 3 files changed, 45 insertions(+)

diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD
index 55e9c98c..70281271 100644
--- a/mediapipe/gpu/BUILD
+++ b/mediapipe/gpu/BUILD
@@ -441,6 +441,21 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "gpu_buffer_storage_yuv_image",
+    srcs = ["gpu_buffer_storage_yuv_image.cc"],
+    hdrs = ["gpu_buffer_storage_yuv_image.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":gpu_buffer_format",
+        ":gpu_buffer_storage",
+        "//mediapipe/framework/formats:yuv_image",
+        "//third_party/libyuv",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
 cc_library(
     name = "gpu_buffer_storage_ahwb",
     srcs = ["gpu_buffer_storage_ahwb.cc"],
@@ -1187,3 +1202,17 @@ mediapipe_cc_test(
         "//mediapipe/framework/port:gtest_main",
     ],
 )
+
+mediapipe_cc_test(
+    name = "gpu_buffer_storage_yuv_image_test",
+    size = "small",
+    srcs = ["gpu_buffer_storage_yuv_image_test.cc"],
+    exclude_platforms = [
+        "ios",
+    ],
+    deps = [
+        ":gpu_buffer_format",
+        ":gpu_buffer_storage_yuv_image",
+        "//mediapipe/framework/port:gtest_main",
+    ],
+)
diff --git a/mediapipe/gpu/gpu_buffer_format.cc b/mediapipe/gpu/gpu_buffer_format.cc
index 1dcd58e6..8e2e3858 100644
--- a/mediapipe/gpu/gpu_buffer_format.cc
+++ b/mediapipe/gpu/gpu_buffer_format.cc
@@ -212,6 +212,10 @@ ImageFormat::Format ImageFormatForGpuBufferFormat(GpuBufferFormat format) {
     case GpuBufferFormat::kTwoComponentHalf16:
     case GpuBufferFormat::kRGBAHalf64:
     case GpuBufferFormat::kRGBAFloat128:
+    case GpuBufferFormat::kNV12:
+    case GpuBufferFormat::kNV21:
+    case GpuBufferFormat::kI420:
+    case GpuBufferFormat::kYV12:
     case GpuBufferFormat::kUnknown:
       return ImageFormat::UNKNOWN;
   }
diff --git a/mediapipe/gpu/gpu_buffer_format.h b/mediapipe/gpu/gpu_buffer_format.h
index 06c5a043..5d77afeb 100644
--- a/mediapipe/gpu/gpu_buffer_format.h
+++ b/mediapipe/gpu/gpu_buffer_format.h
@@ -52,6 +52,14 @@ enum class GpuBufferFormat : uint32_t {
   kRGB24 = 0x00000018,  // Note: prefer BGRA32 whenever possible.
   kRGBAHalf64 = MEDIAPIPE_FOURCC('R', 'G', 'h', 'A'),
   kRGBAFloat128 = MEDIAPIPE_FOURCC('R', 'G', 'f', 'A'),
+  // 8-bit Y plane + interleaved 8-bit U/V plane with 2x2 subsampling.
+  kNV12 = MEDIAPIPE_FOURCC('N', 'V', '1', '2'),
+  // 8-bit Y plane + interleaved 8-bit V/U plane with 2x2 subsampling.
+  kNV21 = MEDIAPIPE_FOURCC('N', 'V', '2', '1'),
+  // 8-bit Y plane + non-interleaved 8-bit U/V planes with 2x2 subsampling.
+  kI420 = MEDIAPIPE_FOURCC('I', '4', '2', '0'),
+  // 8-bit Y plane + non-interleaved 8-bit V/U planes with 2x2 subsampling.
+  kYV12 = MEDIAPIPE_FOURCC('Y', 'V', '1', '2'),
 };
 
 #if !MEDIAPIPE_DISABLE_GPU
@@ -111,6 +119,10 @@ inline OSType CVPixelFormatForGpuBufferFormat(GpuBufferFormat format) {
       return kCVPixelFormatType_64RGBAHalf;
     case GpuBufferFormat::kRGBAFloat128:
       return kCVPixelFormatType_128RGBAFloat;
+    case GpuBufferFormat::kNV12:
+    case GpuBufferFormat::kNV21:
+    case GpuBufferFormat::kI420:
+    case GpuBufferFormat::kYV12:
     case GpuBufferFormat::kUnknown:
       return -1;
   }