Add no-copy Image getter for JNIS

PiperOrigin-RevId: 519909198

mediapipe/java/com/google/mediapipe/framework/PacketGetter.java

@@ -24,6 +24,7 @@ import com.google.protobuf.Parser;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
+import javax.annotation.Nullable;
 
 /**
  * Converts the {@link Packet} to java accessible data types.
@@ -187,6 +188,10 @@ public final class PacketGetter {
     return nativeGetImageHeight(packet.getNativeHandle());
   }
 
+  public static int getImageNumChannels(final Packet packet) {
+    return nativeGetImageNumChannels(packet.getNativeHandle());
+  }
+
   /**
    * Returns the native image buffer in ByteBuffer. It assumes the output buffer stores pixels
    * contiguously. It returns false if this assumption does not hold.
@@ -199,6 +204,18 @@ public final class PacketGetter {
     return nativeGetImageData(packet.getNativeHandle(), buffer);
   }
 
+  /**
+   * Returns a read-only view of the native image buffer as a ByteBuffer. As this method does not
+   * copy the data, the result only remains valid while the backing MediaPipe image is on the stack.
+   * The image must store contiguous pixels, otherwise the method returns {@code null}.
+   *
+   * <p>Note: this function does not assume the pixel format.
+   */
+  @Nullable
+  public static ByteBuffer getImageDataDirectly(final Packet packet) {
+    return nativeGetImageDataDirect(packet.getNativeHandle()).asReadOnlyBuffer();
+  }
+
   /** Returns the size of Image list. This helps to determine size of allocated ByteBuffer array. */
   public static int getImageListSize(final Packet packet) {
     return nativeGetImageListSize(packet.getNativeHandle());
@@ -384,8 +401,12 @@ public final class PacketGetter {
 
   private static native int nativeGetImageHeight(long nativePacketHandle);
 
+  private static native int nativeGetImageNumChannels(long nativePacketHandle);
+
   private static native boolean nativeGetImageData(long nativePacketHandle, ByteBuffer buffer);
 
+  private static native ByteBuffer nativeGetImageDataDirect(long nativePacketHandle);
+
   private static native int nativeGetImageListSize(long nativePacketHandle);
 
   private static native boolean nativeGetImageList(

mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.cc

@@ -334,6 +334,20 @@ JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageHeight)(
   return image.Height();
 }
 
+JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageNumChannels)(
+    JNIEnv* env, jobject thiz, jlong packet) {
+  mediapipe::Packet mediapipe_packet =
+      mediapipe::android::Graph::GetPacketFromHandle(packet);
+  const bool is_image =
+      mediapipe_packet.ValidateAsType<mediapipe::Image>().ok();
+  const mediapipe::ImageFrame& image =
+      is_image ? *GetFromNativeHandle<mediapipe::Image>(packet)
+                      .GetImageFrameSharedPtr()
+                      .get()
+               : GetFromNativeHandle<mediapipe::ImageFrame>(packet);
+  return image.NumberOfChannels();
+}
+
 JNIEXPORT jboolean JNICALL PACKET_GETTER_METHOD(nativeGetImageData)(
     JNIEnv* env, jobject thiz, jlong packet, jobject byte_buffer) {
   mediapipe::Packet mediapipe_packet =
@@ -348,6 +362,31 @@ JNIEXPORT jboolean JNICALL PACKET_GETTER_METHOD(nativeGetImageData)(
   return CopyImageDataToByteBuffer(env, image, byte_buffer);
 }
 
+JNIEXPORT jobject JNICALL PACKET_GETTER_METHOD(nativeGetImageDataDirect)(
+    JNIEnv* env, jobject thiz, jlong packet) {
+  mediapipe::Packet mediapipe_packet =
+      mediapipe::android::Graph::GetPacketFromHandle(packet);
+  const bool is_image =
+      mediapipe_packet.ValidateAsType<mediapipe::Image>().ok();
+  const mediapipe::ImageFrame& image =
+      is_image ? *GetFromNativeHandle<mediapipe::Image>(packet)
+                      .GetImageFrameSharedPtr()
+                      .get()
+               : GetFromNativeHandle<mediapipe::ImageFrame>(packet);
+  if (!image.IsContiguous()) {
+    return NULL;
+  }
+
+  // We need to get a mutable data pointer to create a ByteBuffer in Java. Since
+  // we are returning a read-only ByteBuffer via the API, we essentially retain
+  // the original const qualifier.
+  mediapipe::ImageFrame& mutable_image =
+      const_cast<mediapipe::ImageFrame&>(image);
+  void* unsafe_ptr = static_cast<void*>(mutable_image.MutablePixelData());
+  return env->NewDirectByteBuffer(unsafe_ptr,
+                                  image.PixelDataSizeStoredContiguously());
+}
+
 JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageListSize)(
     JNIEnv* env, jobject thiz, jlong packet) {
   const auto& image_list =

mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.h

@@ -101,11 +101,18 @@ JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageHeight)(JNIEnv* env,
                                                                   jobject thiz,
                                                                   jlong packet);
 
+JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageNumChannels)(
+    JNIEnv* env, jobject thiz, jlong packet);
+
 // Before calling this, the byte_buffer needs to have the correct allocated
 // size.
 JNIEXPORT jboolean JNICALL PACKET_GETTER_METHOD(nativeGetImageData)(
     JNIEnv* env, jobject thiz, jlong packet, jobject byte_buffer);
 
+// Returns the existing byte buffer of a MediaPipe image.
+JNIEXPORT jobject JNICALL PACKET_GETTER_METHOD(nativeGetImageDataDirect)(
+    JNIEnv* env, jobject thiz, jlong packet);
+
 // Return the vector size of std::vector<Image>.
 JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageListSize)(
     JNIEnv* env, jobject thiz, jlong packet);