[apriltag] Make AprilTagDetector.detect() use RawFrame instead of Ope…

…nCV Mat
wpilibsuite · Nov 9, 2024 · 5dd380b · 5dd380b
1 parent 280d2c7
commit 5dd380b
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 9 deletions.
diff --git a/apriltag/src/main/java/edu/wpi/first/apriltag/AprilTagDetector.java b/apriltag/src/main/java/edu/wpi/first/apriltag/AprilTagDetector.java
@@ -5,7 +5,7 @@
 package edu.wpi.first.apriltag;
 
 import edu.wpi.first.apriltag.jni.AprilTagJNI;
-import org.opencv.core.Mat;
+import edu.wpi.first.util.RawFrame;
 
 /**
  * An AprilTag detector engine. This is expensive to set up and tear down, so most use cases should
@@ -293,11 +293,12 @@ public void clearFamilies() {
    *
    * <p>The image must be grayscale.
    *
-   * @param img 8-bit OpenCV Mat image
+   * @param frame The frame object containing an 8-bit image.
    * @return Results (array of AprilTagDetection)
    */
-  public AprilTagDetection[] detect(Mat img) {
-    return AprilTagJNI.detect(m_native, img.cols(), img.rows(), img.cols(), img.dataAddr());
+  public AprilTagDetection[] detect(RawFrame frame) {
+    return AprilTagJNI.detect(
+        m_native, frame.getWidth(), frame.getHeight(), frame.getWidth(), frame.getDataPtr());
   }
 
   private long m_native;

diff --git a/apriltag/src/test/java/edu/wpi/first/apriltag/AprilTagDetectorTest.java b/apriltag/src/test/java/edu/wpi/first/apriltag/AprilTagDetectorTest.java
@@ -10,6 +10,8 @@
 
 import edu.wpi.first.math.geometry.Transform3d;
 import edu.wpi.first.math.util.Units;
+import edu.wpi.first.util.PixelFormat;
+import edu.wpi.first.util.RawFrame;
 import edu.wpi.first.util.RuntimeLoader;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -143,8 +145,15 @@ void testDecodeAndPose() {
       fail(ex);
       return;
     }
+
+    var frameBytes = new byte[image.width() * image.height()];
+    image.get(0, 0, frameBytes);
+
+    var frame = new RawFrame();
+    frame.setData(frameBytes, image.width(), image.height(), 1, PixelFormat.kGray);
+
     try {
-      AprilTagDetection[] results = detector.detect(image);
+      AprilTagDetection[] results = detector.detect(frame);
       assertEquals(1, results.length);
       assertEquals("tag36h11", results[0].getFamily());
       assertEquals(1, results[0].getId());
@@ -176,8 +185,15 @@ void testPoseRotatedX() {
       fail(ex);
       return;
     }
+
+    var frameBytes = new byte[image.width() * image.height()];
+    image.get(0, 0, frameBytes);
+
+    var frame = new RawFrame();
+    frame.setData(frameBytes, image.width(), image.height(), 1, PixelFormat.kGray);
+
     try {
-      AprilTagDetection[] results = detector.detect(image);
+      AprilTagDetection[] results = detector.detect(frame);
       assertEquals(1, results.length);
 
       var estimator =
@@ -209,8 +225,15 @@ void testPoseRotatedY() {
       fail(ex);
       return;
     }
+
+    var frameBytes = new byte[image.width() * image.height()];
+    image.get(0, 0, frameBytes);
+
+    var frame = new RawFrame();
+    frame.setData(frameBytes, image.width(), image.height(), 1, PixelFormat.kGray);
+
     try {
-      AprilTagDetection[] results = detector.detect(image);
+      AprilTagDetection[] results = detector.detect(frame);
       assertEquals(1, results.length);
 
       var estimator =
@@ -239,8 +262,15 @@ void testPoseStraightOn() {
       fail(ex);
       return;
     }
+
+    var frameBytes = new byte[image.width() * image.height()];
+    image.get(0, 0, frameBytes);
+
+    var frame = new RawFrame();
+    frame.setData(frameBytes, image.width(), image.height(), 1, PixelFormat.kGray);
+
     try {
-      AprilTagDetection[] results = detector.detect(image);
+      AprilTagDetection[] results = detector.detect(frame);
       assertEquals(1, results.length);
 
       var estimator =

diff --git a/wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/examples/apriltagsvision/Robot.java b/wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/examples/apriltagsvision/Robot.java
@@ -16,6 +16,8 @@
 import edu.wpi.first.networktables.IntegerArrayPublisher;
 import edu.wpi.first.networktables.NetworkTable;
 import edu.wpi.first.networktables.NetworkTableInstance;
+import edu.wpi.first.util.PixelFormat;
+import edu.wpi.first.util.RawFrame;
 import edu.wpi.first.wpilibj.TimedRobot;
 import java.util.ArrayList;
 import org.opencv.core.Mat;
@@ -65,6 +67,9 @@ void apriltagVisionThreadProc() {
     var mat = new Mat();
     var grayMat = new Mat();
 
+    var frameBytes = new byte[640 * 480];
+    var frame = new RawFrame();
+
     // Instantiate once
     ArrayList<Long> tags = new ArrayList<>();
     var outlineColor = new Scalar(0, 255, 0);
@@ -89,7 +94,10 @@ void apriltagVisionThreadProc() {
 
       Imgproc.cvtColor(mat, grayMat, Imgproc.COLOR_RGB2GRAY);
 
-      AprilTagDetection[] detections = detector.detect(grayMat);
+      grayMat.get(0, 0, frameBytes);
+      frame.setData(frameBytes, grayMat.width(), grayMat.height(), 1, PixelFormat.kGray);
+
+      AprilTagDetection[] detections = detector.detect(frame);
 
       // have not seen any tags yet
       tags.clear();

diff --git a/wpiutil/src/main/java/edu/wpi/first/util/RawFrame.java b/wpiutil/src/main/java/edu/wpi/first/util/RawFrame.java
@@ -88,6 +88,22 @@ public void setData(ByteBuffer data, int width, int height, int stride, PixelFor
         m_nativeObj, data, data.limit(), width, height, stride, pixelFormat.getValue());
   }
 
+  /**
+   * Set frame data.
+   *
+   * @param data A byte array containing the frame data.
+   * @param width The width of the frame, in pixels
+   * @param height The height of the frame, in pixels
+   * @param stride The number of bytes in each row of image data
+   * @param pixelFormat The PixelFormat of the frame
+   */
+  public void setData(byte[] data, int width, int height, int stride, PixelFormat pixelFormat) {
+    var dataByteBuffer = ByteBuffer.allocateDirect(width * height * stride);
+    dataByteBuffer.put(data);
+
+    setData(dataByteBuffer, width, height, stride, pixelFormat);
+  }
+
   /**
    * Call to set frame information.
    *