Add no-copy Image getter for JNIS

PiperOrigin-RevId: 519909198
This commit is contained in:
Sebastian Schmidt 2023-03-27 20:57:04 -07:00 committed by Copybara-Service
parent 59b3150fff
commit 105e7b7467
3 changed files with 67 additions and 0 deletions

View File

@ -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(

View File

@ -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 =

View File

@ -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);