[apriltag] Fix Java generation functions (#6063)

wpilibsuite · Dec 22, 2023 · bcef6c5 · bcef6c5
1 parent 4059e0c
commit bcef6c5
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 19 deletions.
diff --git a/apriltag/src/main/java/edu/wpi/first/apriltag/AprilTag.java b/apriltag/src/main/java/edu/wpi/first/apriltag/AprilTag.java
@@ -54,9 +54,7 @@ public String toString() {
    * @return A RawFrame containing the AprilTag image
    */
   public static RawFrame generate16h5AprilTagImage(int id) {
-    RawFrame generatedImage = new RawFrame();
-    AprilTagJNI.generate16h5AprilTagImage(id, generatedImage.getDataPtr());
-    return generatedImage;
+    return AprilTagJNI.generate16h5AprilTagImage(id);
   }
 
   /**
@@ -66,8 +64,6 @@ public static RawFrame generate16h5AprilTagImage(int id) {
    * @return A RawFrame containing the AprilTag image
    */
   public static RawFrame generate36h11AprilTagImage(int id) {
-    RawFrame generatedImage = new RawFrame();
-    AprilTagJNI.generate36h11AprilTagImage(id, generatedImage.getDataPtr());
-    return generatedImage;
+    return AprilTagJNI.generate36h11AprilTagImage(id);
   }
 }
diff --git a/apriltag/src/main/java/edu/wpi/first/apriltag/jni/AprilTagJNI.java b/apriltag/src/main/java/edu/wpi/first/apriltag/jni/AprilTagJNI.java
@@ -8,6 +8,7 @@
 import edu.wpi.first.apriltag.AprilTagDetector;
 import edu.wpi.first.apriltag.AprilTagPoseEstimate;
 import edu.wpi.first.math.geometry.Transform3d;
+import edu.wpi.first.util.RawFrame;
 import edu.wpi.first.util.RuntimeLoader;
 import java.io.IOException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -190,7 +191,7 @@ public static native Transform3d estimatePose(
       double cx,
       double cy);
 
-  public static native void generate16h5AprilTagImage(int id, long nativeAddr);
+  public static native RawFrame generate16h5AprilTagImage(int id);
 
-  public static native void generate36h11AprilTagImage(int id, long nativeAddr);
+  public static native RawFrame generate36h11AprilTagImage(int id);
 }
diff --git a/apriltag/src/main/native/cpp/jni/AprilTagJNI.cpp b/apriltag/src/main/native/cpp/jni/AprilTagJNI.cpp
@@ -5,6 +5,7 @@
 #include <cstdio>
 #include <cstring>
 
+#include <wpi/RawFrame.h>
 #include <wpi/jni_util.h>
 
 #include "edu_wpi_first_apriltag_jni_AprilTagJNI.h"
@@ -25,6 +26,7 @@ static JClass quaternionCls;
 static JClass rotation3dCls;
 static JClass transform3dCls;
 static JClass translation3dCls;
+static JClass rawFrameCls;
 static JException illegalArgEx;
 static JException nullPointerEx;
 
@@ -37,7 +39,8 @@ static const JClassInit classes[] = {
     {"edu/wpi/first/math/geometry/Quaternion", &quaternionCls},
     {"edu/wpi/first/math/geometry/Rotation3d", &rotation3dCls},
     {"edu/wpi/first/math/geometry/Transform3d", &transform3dCls},
-    {"edu/wpi/first/math/geometry/Translation3d", &translation3dCls}};
+    {"edu/wpi/first/math/geometry/Translation3d", &translation3dCls},
+    {"edu/wpi/first/util/RawFrame", &rawFrameCls}};
 
 static const JExceptionInit exceptions[] = {
     {"java/lang/IllegalArgumentException", &illegalArgEx},
@@ -316,6 +319,21 @@ static jobject MakeJObject(JNIEnv* env, const AprilTagPoseEstimate& est) {
                         static_cast<jdouble>(est.error2));
 }
 
+static jobject MakeJObject(JNIEnv* env, const wpi::RawFrame& frame) {
+  static jmethodID constructor = env->GetMethodID(rawFrameCls, "<init>", "()V");
+
+  static jmethodID setData =
+      env->GetMethodID(rawFrameCls, "setData", "(Ljava/nio/ByteBuffer;JIIII)V");
+
+  jobject retVal = env->NewObject(rawFrameCls, constructor);
+  env->CallVoidMethod(
+      retVal, setData, env->NewDirectByteBuffer(frame.data, frame.totalData),
+      static_cast<jlong>(reinterpret_cast<intptr_t>(frame.data)),
+      static_cast<jint>(frame.dataLength), static_cast<jint>(frame.width),
+      static_cast<jint>(frame.height), static_cast<jint>(frame.pixelFormat));
+  return retVal;
+}
+
 extern "C" {
 
 /*
@@ -591,27 +609,25 @@ Java_edu_wpi_first_apriltag_jni_AprilTagJNI_estimatePose
 /*
  * Class:     edu_wpi_first_apriltag_jni_AprilTagJNI
  * Method:    generate16h5AprilTagImage
- * Signature: (IJ)V
+ * Signature: (I)Ljava/lang/Object;
  */
-JNIEXPORT void JNICALL
+JNIEXPORT jobject JNICALL
 Java_edu_wpi_first_apriltag_jni_AprilTagJNI_generate16h5AprilTagImage
-  (JNIEnv* env, jclass, jint id, jlong framePtr)
+  (JNIEnv* env, jclass, jint id)
 {
-  wpi::RawFrame* javaRawFrame = (wpi::RawFrame*)framePtr;
-  *javaRawFrame = AprilTag::Generate16h5AprilTagImage(id);
+  return MakeJObject(env, AprilTag::Generate16h5AprilTagImage(id));
 }
 
 /*
  * Class:     edu_wpi_first_apriltag_jni_AprilTagJNI
  * Method:    generate36h11AprilTagImage
- * Signature: (IJ)V
+ * Signature: (I)Ljava/lang/Object;
  */
-JNIEXPORT void JNICALL
+JNIEXPORT jobject JNICALL
 Java_edu_wpi_first_apriltag_jni_AprilTagJNI_generate36h11AprilTagImage
-  (JNIEnv* env, jclass, jint id, jlong framePtr)
+  (JNIEnv* env, jclass, jint id)
 {
-  wpi::RawFrame* javaRawFrame = (wpi::RawFrame*)framePtr;
-  *javaRawFrame = AprilTag::Generate36h11AprilTagImage(id);
+  return MakeJObject(env, AprilTag::Generate36h11AprilTagImage(id));
 }
 
 }  // extern "C"