Add jni functions to get image size from a list of image

PiperOrigin-RevId: 524080111

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

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<mediapipe::Image>& 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<std::vector<mediapipe::Image>>(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<std::vector<mediapipe::Image>>(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) {

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<Image>.
+JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageWidthFromImageList)(
+    JNIEnv* env, jobject thiz, jlong packet);
+
+// Return the image height of first image from std::vector<Image>.
+JNIEXPORT jint JNICALL PACKET_GETTER_METHOD(nativeGetImageHeightFromImageList)(
+    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);