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

PiperOrigin-RevId: 524080111
This commit is contained in:
MediaPipe Team 2023-04-13 13:03:29 -07:00 committed by Copybara-Service
parent ad169b9123
commit a59b29ea24
3 changed files with 71 additions and 0 deletions

View File

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

View File

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

View File

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