diff --git a/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java b/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java index 5a13e9c05..1c1daadcc 100644 --- a/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java +++ b/mediapipe/java/com/google/mediapipe/framework/PacketGetter.java @@ -221,6 +221,22 @@ public final class PacketGetter { return nativeGetImageListSize(packet.getNativeHandle()); } + /** + * Returns the width of first image in an image list. This helps to determine size of allocated + * ByteBuffer array. + */ + public static int getImageWidthFromImageList(final Packet packet) { + return nativeGetImageWidthFromImageList(packet.getNativeHandle()); + } + + /** + * Returns the height of first image in an image list. This helps to determine size of allocated + * ByteBuffer array. + */ + public static int getImageHeightFromImageList(final Packet packet) { + return nativeGetImageHeightFromImageList(packet.getNativeHandle()); + } + /** * Assign the native image buffer array in given ByteBuffer array. It assumes given ByteBuffer * array has the the same size of image list packet, and assumes the output buffer stores pixels @@ -409,6 +425,10 @@ public final class PacketGetter { private static native int nativeGetImageListSize(long nativePacketHandle); + private static native int nativeGetImageWidthFromImageList(long nativePacketHandle); + + private static native int nativeGetImageHeightFromImageList(long nativePacketHandle); + private static native boolean nativeGetImageList( long nativePacketHandle, ByteBuffer[] bufferArray, boolean deepCopy); 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 cc273bca4..093b147a2 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 @@ -16,6 +16,7 @@ #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" #include "mediapipe/framework/calculator.pb.h" #include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/image_frame.h" @@ -86,6 +87,22 @@ bool CopyImageDataToByteBuffer(JNIEnv* env, const mediapipe::ImageFrame& image, return true; } +void CheckImageSizeInImageList(JNIEnv* env, + const std::vector& image_list, + int height, int width, int channels) { + for (int i = 0; i < image_list.size(); ++i) { + if (image_list[i].height() != height || image_list[i].width() != width || + image_list[i].channels() != channels) { + ThrowIfError(env, absl::InvalidArgumentError(absl::StrFormat( + "Expect images in the image list having the same " + "size: (%d, %d, %d), but get image at index %d " + "with size: (%d, %d, %d)", + height, width, channels, i, image_list[i].height(), + image_list[i].width(), image_list[i].channels()))); + } + } +} + } // namespace JNIEXPORT jlong JNICALL PACKET_GETTER_METHOD(nativeGetPacketFromReference)( @@ -394,6 +411,32 @@ JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageListSize)( return image_list.size(); } +JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageWidthFromImageList)( + JNIEnv* env, jobject thiz, jlong packet) { + const auto& image_list = + GetFromNativeHandle>(packet); + if (image_list.empty()) { + ThrowIfError(env, absl::InvalidArgumentError( + "Image list from the packet is empty.")); + } + CheckImageSizeInImageList(env, image_list, image_list[0].height(), + image_list[0].width(), image_list[0].channels()); + return image_list[0].width(); +} + +JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageHeightFromImageList)( + JNIEnv* env, jobject thiz, jlong packet) { + const auto& image_list = + GetFromNativeHandle>(packet); + if (image_list.empty()) { + ThrowIfError(env, absl::InvalidArgumentError( + "Image list from the packet is empty.")); + } + CheckImageSizeInImageList(env, image_list, image_list[0].height(), + image_list[0].width(), image_list[0].channels()); + return image_list[0].height(); +} + JNIEXPORT jboolean JNICALL PACKET_GETTER_METHOD(nativeGetImageList)( JNIEnv* env, jobject thiz, jlong packet, jobjectArray byte_buffer_array, jboolean deep_copy) { 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 e35ea5243..202795307 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 @@ -113,6 +113,14 @@ JNIEXPORT jboolean JNICALL PACKET_GETTER_METHOD(nativeGetImageData)( JNIEXPORT jobject JNICALL PACKET_GETTER_METHOD(nativeGetImageDataDirect)( JNIEnv* env, jobject thiz, jlong packet); +// Return the image width of first image from std::vector. +JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageWidthFromImageList)( + JNIEnv* env, jobject thiz, jlong packet); + +// Return the image height of first image from std::vector. +JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageHeightFromImageList)( + 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);