fix: convert Android420 to I420 then to RGBA

bboettcher3 · Dec 3, 2020 · 4e1caaa · 4e1caaa
1 parent f94a22d
commit 4e1caaa
Show file tree

Hide file tree

Showing 8 changed files with 80 additions and 11 deletions.
diff --git a/mediapipe/java/com/google/mediapipe/framework/Graph.java b/mediapipe/java/com/google/mediapipe/framework/Graph.java
@@ -381,12 +381,12 @@ public synchronized void addConsumablePacketToInputStream(
    */
   public synchronized void sendInputYuvFrame(
       long timestamp, ByteBuffer yBuffer, ByteBuffer uBuffer, ByteBuffer vBuffer, int yStride,
-      int uStride, int vStride, int width, int height, int lensRotation, boolean shouldFlipX, boolean shouldSendLens) {
+      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, yStride, uStride, vStride, 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 yStride, int uStride, int vStride, 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
@@ -154,6 +154,17 @@ public Packet createRgbaImageFrame(ByteBuffer buffer, int width, int height) {
         nativeCreateRgbaImageFrame(mediapipeGraph.getNativeHandle(), buffer, width, height));
   }
 
+  /**
+   * 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 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, yStride, uStride, vStride, uvPixelStride, yDestBuffer, uDestBuffer, vDestBuffer, width, height));
+  }
+
   /**
    * Creates a 1 channel float ImageFrame packet.
    *
@@ -349,6 +360,9 @@ 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 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 y_stride, jint u_stride, jint v_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, 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 y_stride, jint u_stride, jint v_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
@@ -182,6 +182,29 @@ JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateRgbaImageFrame)(
   return CreatePacketWithContext(context, packet);
 }
 
+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 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, 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);
+}
+
 static mediapipe::Packet createAudioPacket(const uint8_t* audio_sample,
                                            int num_samples, int num_channels) {
   std::unique_ptr<::mediapipe::Matrix> matrix(

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
@@ -39,6 +39,12 @@ JNIEXPORT jlong JNICALL PACKET_CREATOR_METHOD(nativeCreateRgbaImageFrame)(
     JNIEnv* env, jobject thiz, jlong context, jobject byte_buffer, jint width,
     jint height);
 
+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 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)(
     JNIEnv* env, jobject thiz, jlong context, jobject byte_buffer, jint width,
     jint height);

diff --git a/mediapipe/util/image_frame_util.cc b/mediapipe/util/image_frame_util.cc
@@ -208,18 +208,33 @@ void YUVToRgbaImageFrame(
         int yStride,
         int uStride,
         int vStride,
+        int uvPixelStride,
+        uint8* destYData,
+        uint8* destUData,
+        uint8* destVData,
         int width,
         int height,
         ImageFrame* image_frame
 ) {
   CHECK(image_frame);
   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(yData, yStride,  //
-                         uData, uStride,  //
-                         vData, vStride,  //
+  rv = libyuv::I420ToABGR(destYData, yStride,  //
+                         destUData, uStride,  //
+                         destVData, vStride,  //
                          image_frame->MutablePixelData(),
                          image_frame->WidthStep(), width, height);
   CHECK_EQ(0, rv);

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 yStride, int uStride, int vStride, 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