From 105e7b74677b9f0923edc7f9c197e158df8051f6 Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Mon, 27 Mar 2023 20:57:04 -0700 Subject: [PATCH] Add no-copy Image getter for JNIS PiperOrigin-RevId: 519909198 --- .../mediapipe/framework/PacketGetter.java | 21 ++++++++++ .../framework/jni/packet_getter_jni.cc | 39 +++++++++++++++++++ .../framework/jni/packet_getter_jni.h | 7 ++++ 3 files changed, 67 insertions(+) diff --git a/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java b/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java index 92cf723e..5a13e9c0 100644 --- a/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java +++ b/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}. + * + *

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( diff --git a/mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.cc b/mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.cc index 234209b8..d5bd773f 100644 --- a/mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.cc +++ b/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().ok(); + const mediapipe::ImageFrame& image = + is_image ? *GetFromNativeHandle(packet) + .GetImageFrameSharedPtr() + .get() + : GetFromNativeHandle(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().ok(); + const mediapipe::ImageFrame& image = + is_image ? *GetFromNativeHandle(packet) + .GetImageFrameSharedPtr() + .get() + : GetFromNativeHandle(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(image); + void* unsafe_ptr = static_cast(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 = diff --git a/mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.h b/mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.h index 4602ebd5..e35ea524 100644 --- a/mediapipe/java/com/google/mediapipe/framework/jni/packet_getter_jni.h +++ b/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. JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageListSize)( JNIEnv* env, jobject thiz, jlong packet);