Merge pull request google-ai-edge#12 from jianinz/fix/improve-yuv-rgb…

…a-conversion-effeciency perf: improve YUV to RGBA conversion efficiency
bboettcher3 · Dec 3, 2020 · c27234f · c27234f
2 parents bf60e5f + 4e1caaa
commit c27234f
Show file tree

Hide file tree

Showing 10 changed files with 72 additions and 30 deletions.
diff --git a/mediapipe/calculators/image/image_transformation_calculator.cc b/mediapipe/calculators/image/image_transformation_calculator.cc
@@ -21,7 +21,6 @@
 #include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/framework/port/status.h"
 #include "mediapipe/gpu/scale_mode.pb.h"
-
 #if !defined(MEDIAPIPE_DISABLE_GPU)
 #include "mediapipe/gpu/gl_calculator_helper.h"
 #include "mediapipe/gpu/gl_quad_renderer.h"

diff --git a/mediapipe/gpu/gpu_buffer_to_image_frame_calculator.cc b/mediapipe/gpu/gpu_buffer_to_image_frame_calculator.cc
@@ -35,7 +35,6 @@ class GpuBufferToImageFrameCalculator : public CalculatorBase {
 
   ::mediapipe::Status Open(CalculatorContext* cc) override;
   ::mediapipe::Status Process(CalculatorContext* cc) override;
-
  private:
 #if !MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
   GlCalculatorHelper helper_;
@@ -98,7 +97,6 @@ ::mediapipe::Status GpuBufferToImageFrameCalculator::Process(
     return ::mediapipe::OkStatus();
   }
 #endif  // defined(HAVE_GPU_BUFFER)
-
   return ::mediapipe::Status(::mediapipe::StatusCode::kInvalidArgument,
                              "Input packets must be ImageFrame or GpuBuffer.");
 }

diff --git a/mediapipe/java/com/google/mediapipe/framework/Graph.java b/mediapipe/java/com/google/mediapipe/framework/Graph.java
@@ -380,13 +380,13 @@ public synchronized void addConsumablePacketToInputStream(
    * @throws MediaPipeException for any error status.
    */
   public synchronized void sendInputYuvFrame(
-      long timestamp, ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer,
-      int uvStride, int width, int height, int lensRotation, boolean shouldFlipX, boolean shouldSendLens) {
+      long timestamp, ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int yStride,
+      int uStride, int vStride, int uvPixelStride, ByteBuffer yDestBuffer, ByteBuffer uDestBuffer, ByteBuffer vDestBuffer, int width, int height, int lensRotation, boolean shouldFlipX, boolean shouldSendLens) {
     Preconditions.checkState(
         nativeGraphHandle != 0, "Invalid context, tearDown() might have been called.");
     if (graphRunning) {
       // We perform the packet creation and sending all natively for performance
-      nativeSendInputYuvFrame(nativeGraphHandle, timestamp, yBuffer, uBuffer, vBuffer, uvStride, width, height,
+      nativeSendInputYuvFrame(nativeGraphHandle, timestamp, yBuffer, uBuffer, vBuffer, yStride, uStride, vStride, uvPixelStride, yDestBuffer, uDestBuffer, vDestBuffer, width, height,
           lensRotation, shouldFlipX, shouldSendLens);
     }
   }
@@ -636,7 +636,9 @@ private native void nativeMovePacketToInputStream(
       long context, String streamName, long packet, long timestamp);
 
   private native void nativeSendInputYuvFrame(long context, long timestamp,
-      ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int uvStride, int width, int height,
+      ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer,
+      int yStride, int uStride, int vStride, int uvPixelStride, ByteBuffer yDestBuffer,
+      ByteBuffer uDestBuffer, ByteBuffer vDestBuffer, int width, int height,
       int lensRotation, boolean shouldFlipX, boolean shouldSendLens);
 
   private native void nativeSetGraphInputStreamBlockingMode(long context, boolean mode);

diff --git a/mediapipe/java/com/google/mediapipe/framework/PacketCreator.java b/mediapipe/java/com/google/mediapipe/framework/PacketCreator.java
@@ -155,14 +155,14 @@ public Packet createRgbaImageFrame(ByteBuffer buffer, int width, int height) {
   }
 
   /**
-   * Creates a 4 channel RGBA ImageFrame packet from a YUV_420_888 buffer.
+   * Creates a 4 channel RGBA ImageFrame packet from a YUV_420_888 buffer. Can be used for debugging so that Android client could view the packet bitmap
    */
-  public Packet createYuvImageFrame(ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int uvStride, int width, int height) {
+  public Packet createYuvImageFrame(ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int yStride, int uStride, int vStride, int uvPixelStride, ByteBuffer yDestBuffer, ByteBuffer uDestBuffer, ByteBuffer vDestBuffer, int width, int height) {
     if (yBuffer.capacity() != width * height) {
       throw new RuntimeException("buffer doesn't have the correct size.");
     }
     return Packet.create(
-        nativeCreateYuvImageFrame(mediapipeGraph.getNativeHandle(), yBuffer, uBuffer, vBuffer, uvStride, width, height));
+        nativeCreateYuvImageFrame(mediapipeGraph.getNativeHandle(), yBuffer, uBuffer, vBuffer, yStride, uStride, vStride, uvPixelStride, yDestBuffer, uDestBuffer, vDestBuffer, width, height));
   }
 
   /**
@@ -361,7 +361,8 @@ private native long nativeCreateGrayscaleImage(
   private native long nativeCreateRgbaImageFrame(
       long context, ByteBuffer buffer, int width, int height);
   private native long nativeCreateYuvImageFrame(
-      long context, ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int uvStride, int width, int height);
+      long context, ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int yStride, int uStride, int vStride,
+      int uvPixelStride, ByteBuffer yDestBuffer, ByteBuffer uDestBuffer, ByteBuffer vDestBuffer, int width, int height);
   private native long nativeCreateFloatImageFrame(
       long context, FloatBuffer buffer, int width, int height);
   private native long nativeCreateInt16(long context, short value);

diff --git a/mediapipe/java/com/google/mediapipe/framework/jni/graph_jni.cc b/mediapipe/java/com/google/mediapipe/framework/jni/graph_jni.cc
@@ -274,8 +274,9 @@ JNIEXPORT void JNICALL GRAPH_METHOD(nativeMovePacketToInputStream)(
 JNIEXPORT void JNICALL GRAPH_METHOD(nativeSendInputYuvFrame)(
     JNIEnv* env, jobject thiz, jlong context, jlong timestamp,
     jobject y_byte_buffer, jobject u_byte_buffer, jobject v_byte_buffer,
-    jint uv_stride, jint width, jint height,
-    jint lensRotation, jboolean shouldFlipX, jboolean shouldSendLens) {
+    jint y_stride, jint u_stride, jint v_stride, jint uv_PixelStride,
+    jobject y_dest_byte_buffer, jobject u_dest_byte_buffer, jobject v_dest_byte_buffer,
+    jint width, jint height, jint lensRotation, jboolean shouldFlipX, jboolean shouldSendLens) {
   mediapipe::android::Graph* mediapipe_graph =
       reinterpret_cast<mediapipe::android::Graph*>(context);
 
@@ -297,10 +298,16 @@ JNIEXPORT void JNICALL GRAPH_METHOD(nativeSendInputYuvFrame)(
   uint8* y_data = (uint8*)env->GetDirectBufferAddress(y_byte_buffer);
   uint8* u_data = (uint8*)env->GetDirectBufferAddress(u_byte_buffer);
   uint8* v_data = (uint8*)env->GetDirectBufferAddress(v_byte_buffer);
+
+  uint8* y_dest_data = (uint8*)env->GetDirectBufferAddress(y_dest_byte_buffer);
+  uint8* u_dest_data = (uint8*)env->GetDirectBufferAddress(u_dest_byte_buffer);
+  uint8* v_dest_data = (uint8*)env->GetDirectBufferAddress(v_dest_byte_buffer);
+
   auto imageFrame = absl::make_unique<::mediapipe::ImageFrame>(
       mediapipe::ImageFormat::SRGBA, width, height, 8);
   mediapipe::image_frame_util::YUVToRgbaImageFrame(y_data,
-                              u_data, v_data, uv_stride,
+                              u_data, v_data, y_stride, u_stride, v_stride,
+                              uv_PixelStride, y_dest_data, u_dest_data, v_dest_data,
                               width, height, imageFrame.get());
   mediapipe::Packet imagePacket = mediapipe::Adopt(imageFrame.release());
   uint64_t imagePacketHandle = CreatePacketWithContext(context, imagePacket);

diff --git a/mediapipe/java/com/google/mediapipe/framework/jni/graph_jni.h b/mediapipe/java/com/google/mediapipe/framework/jni/graph_jni.h
@@ -85,8 +85,9 @@ JNIEXPORT void JNICALL GRAPH_METHOD(nativeMovePacketToInputStream)(
 JNIEXPORT void JNICALL GRAPH_METHOD(nativeSendInputYuvFrame)(
     JNIEnv* env, jobject thiz, jlong context, jlong timestamp,
     jobject y_byte_buffer, jobject u_byte_buffer, jobject v_byte_buffer,
-    jint uv_stride, jint width, jint height,
-    jint lensRotation, jboolean shouldFlipX, jboolean shouldSendLens);
+    jint y_stride, jint u_stride, jint v_stride, jint uv_PixelStride,
+    jobject y_dest_byte_buffer, jobject u_dest_byte_buffer, jobject v_dest_byte_buffer,
+    jint width, jint height, jint lensRotation, jboolean shouldFlipX, jboolean shouldSendLens);
 
 JNIEXPORT void JNICALL GRAPH_METHOD(nativeSetGraphInputStreamBlockingMode)(
     JNIEnv* env, jobject thiz, jlong context, jboolean mode);

diff --git a/mediapipe/java/com/google/mediapipe/framework/jni/packet_creator_jni.cc b/mediapipe/java/com/google/mediapipe/framework/jni/packet_creator_jni.cc
@@ -184,15 +184,22 @@ JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateRgbaImageFrame)(
 
 JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateYuvImageFrame)(
     JNIEnv* env, jobject thiz, jlong context, jobject y_byte_buffer,
-    jobject u_byte_buffer, jobject v_byte_buffer, jint uv_stride,
+    jobject u_byte_buffer, jobject v_byte_buffer, jint y_stride, jint u_stride, jint v_stride,
+    jint uv_PixelStride, jobject y_dest_byte_buffer, jobject u_dest_byte_buffer, jobject v_dest_byte_buffer,
     jint width, jint height) {
   uint8* y_data = (uint8*)env->GetDirectBufferAddress(y_byte_buffer);
   uint8* u_data = (uint8*)env->GetDirectBufferAddress(u_byte_buffer);
   uint8* v_data = (uint8*)env->GetDirectBufferAddress(v_byte_buffer);
+
+  uint8* y_dest_data = (uint8*)env->GetDirectBufferAddress(y_dest_byte_buffer);
+  uint8* u_dest_data = (uint8*)env->GetDirectBufferAddress(u_dest_byte_buffer);
+  uint8* v_dest_data = (uint8*)env->GetDirectBufferAddress(v_dest_byte_buffer);
+
   auto imageFrame = absl::make_unique<::mediapipe::ImageFrame>(
       mediapipe::ImageFormat::SRGBA, width, height, 8);
   mediapipe::image_frame_util::YUVToRgbaImageFrame(y_data,
-                              u_data, v_data, uv_stride,
+                              u_data, v_data, y_stride, u_stride, v_stride, uv_PixelStride,
+                              y_dest_data, u_dest_data, v_dest_data,
                               width, height, imageFrame.get());
   mediapipe::Packet packet = mediapipe::Adopt(imageFrame.release());
   return CreatePacketWithContext(context, packet);

diff --git a/mediapipe/java/com/google/mediapipe/framework/jni/packet_creator_jni.h b/mediapipe/java/com/google/mediapipe/framework/jni/packet_creator_jni.h
@@ -41,7 +41,8 @@ JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateRgbaImageFrame)(
 
 JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateYuvImageFrame)(
     JNIEnv* env, jobject thiz, jlong context, jobject y_byte_buffer,
-    jobject u_byte_buffer, jobject v_byte_buffer, jint uv_stride,
+    jobject u_byte_buffer, jobject v_byte_buffer, jint y_stride, jint u_stride, jint v_stride,
+    jint uv_PixelStride, jobject y_dest_byte_buffer, jobject u_dest_byte_buffer, jobject v_dest_byte_buffer,
     jint width, jint height);
 
 JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateRgbImageFromRgba)(

diff --git a/mediapipe/util/image_frame_util.cc b/mediapipe/util/image_frame_util.cc
@@ -201,18 +201,43 @@ void YuvToRgbaBuffer(
   }
 }
 
-void YUVToRgbaImageFrame(uint8* yData,
-  uint8* uData,
-  uint8* vData,
-  int uvStride,
-  int width,
-  int height,
-  ImageFrame* image_frame
+void YUVToRgbaImageFrame(
+        uint8* yData,
+        uint8* uData,
+        uint8* vData,
+        int yStride,
+        int uStride,
+        int vStride,
+        int uvPixelStride,
+        uint8* destYData,
+        uint8* destUData,
+        uint8* destVData,
+        int width,
+        int height,
+        ImageFrame* image_frame
 ) {
   CHECK(image_frame);
-  YuvToRgbaBuffer(yData, uData, vData,
-                image_frame->MutablePixelData(),
-                uvStride, width, height);
+  int rv;
+  // Convert Android 420 to I420 first
+  // https://bugs.chromium.org/p/libyuv/issues/detail?id=815&can=1&q=&sort=-id
+  rv = libyuv::Android420ToI420(yData, yStride,
+                                uData, uStride,
+                                vData, vStride,
+                                uvPixelStride,
+                                destYData, yStride,
+                                destUData, uStride,
+                                destVData, vStride,
+                                width, height);
+  CHECK_EQ(0, rv);
+  // libyuv reverses the byte order, so if ABGR is word order, the actual byte order is RBGA, see
+  // https://chromium.googlesource.com/libyuv/libyuv/+/refs/heads/master/source/row_common.cc#544
+  // https://chromium.googlesource.com/libyuv/libyuv/+/refs/heads/master/docs/formats.md#the-argb-fourcc
+  rv = libyuv::I420ToABGR(destYData, yStride,  //
+                         destUData, uStride,  //
+                         destVData, vStride,  //
+                         image_frame->MutablePixelData(),
+                         image_frame->WidthStep(), width, height);
+  CHECK_EQ(0, rv);
 }
 
 void SrgbToMpegYCbCr(const uint8 r, const uint8 g, const uint8 b,  //

diff --git a/mediapipe/util/image_frame_util.h b/mediapipe/util/image_frame_util.h
@@ -66,7 +66,8 @@ void YUVImageToImageFrame(const YUVImage& yuv_image, ImageFrame* image_frame,
 
 // Convert YV planar buffers to an SRGBA ImageFrame.
 void YUVToRgbaImageFrame(uint8* yData, uint8* uData, uint8* vData,
-  int uvStride, int width, int height, ImageFrame* image_frame);
+  int yStride, int uStride, int vStride, int uvPixelStride, uint8* destYData,
+  uint8* destUData, uint8* destVData, int width, int height, ImageFrame* image_frame);
 
 // Convert sRGB values into MPEG YCbCr values.  Notice that MPEG YCbCr
 // values use a smaller range of values than JPEG YCbCr.  The conversion