Uncertainty wrapper

src/main/java/org/photonvision/mrcal/MrCalJNI.java

@@ -19,20 +19,20 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Iterator;
 import java.util.List;
 import org.opencv.core.Point;
+import org.opencv.core.Point3;
 import org.wpilib.math.geometry.Pose3d;
 import org.wpilib.math.geometry.Rotation3d;
 import org.wpilib.math.geometry.Translation3d;
 import org.wpilib.math.linalg.VecBuilder;
 
 public class MrCalJNI {
     public static enum CameraLensModel {
         LENSMODEL_OPENCV5,
         LENSMODEL_OPENCV8,
         LENSMODEL_STEREOGRAPHIC,
         LENSMODEL_SPLINED_STEREOGRAPHIC;
     }
 
     /**
@@ -57,7 +57,7 @@
         }
     }
 
     public static class MrCalResult {
         public boolean success;
         public double[] intrinsics;
         public double rms_error;
@@ -66,7 +66,7 @@
         public double warp_y;
         public int Noutliers;
         public List<Pose3d> optimizedPoses;
         public List<boolean[]> cornersUsed;
 
         public MrCalResult(boolean success) {
             this.success = success;
@@ -279,20 +279,12 @@
      * [here](https://github.com/dkogan/mrcal/blob/7cd9ac4c854a4b244a35f554c9ebd0464d59e9ff/mrcal-calibrate-cameras#L152)
      */
     private static double[] makeObservations(
-            int observationCount,
-            Iterator<MrCalObservation> observationData,
-            int boardWidth,
-            int boardHeight) {
-        double[] observations = new double[boardWidth * boardHeight * 3 * observationCount];
-        Arrays.fill(observations, -1.0);
-
-        for (int b = 0; b < observationCount; b++) {
-            if (!observationData.hasNext()) {
-                // Too few observations
-                return null;
-            }
+            List<MrCalObservation> observations, int boardWidth, int boardHeight) {
+        double[] packedObservations = new double[boardWidth * boardHeight * 3 * observations.size()];
+        Arrays.fill(packedObservations, -1.0);
 
-            final var observation = observationData.next();
+        for (int b = 0; b < observations.size(); b++) {
+            final var observation = observations.get(b);
 
             final var corners = observation.corners;
             final var levels = observation.levels;
@@ -311,9 +303,9 @@
 
                     int i = boardWidth * boardHeight * b + c;
 
-                    observations[i * 3 + 0] = corner.x;
-                    observations[i * 3 + 1] = corner.y;
-                    observations[i * 3 + 2] = level;
+                    packedObservations[i * 3 + 0] = corner.x;
+                    packedObservations[i * 3 + 1] = corner.y;
+                    packedObservations[i * 3 + 2] = level;
                 }
             } else {
                 // Ids present, some corners may be missing
@@ -335,19 +327,38 @@
 
                     int i = boardWidth * boardHeight * b + id;
 
-                    observations[i * 3 + 0] = corner.x;
-                    observations[i * 3 + 1] = corner.y;
-                    observations[i * 3 + 2] = level;
+                    packedObservations[i * 3 + 0] = corner.x;
+                    packedObservations[i * 3 + 1] = corner.y;
+                    packedObservations[i * 3 + 2] = level;
                 }
             }
         }
 
-        if (observationData.hasNext()) {
-            // Too many observations
-            return null;
+        return packedObservations;
+    }
+
+    /**
+     * Convert from WPILib poses to a raw RT array
+     *
+     * @return array of size [numObservations * 6] where each group of 6 is (rvec[0], rvec[1],
+     *     rvec[2], tvec[0], tvec[1], tvec[2]) for the
+     */
+    private static double[] observationPosesToRtRef(List<Pose3d> observationPoses) {
+        double[] ret = new double[observationPoses.size() * 6];
+
+        for (int i = 0; i < observationPoses.size(); i++) {
+            var pose = observationPoses.get(i);
+            var r = pose.getRotation().toVector();
+            var t = pose.getTranslation().toVector();
+            ret[i * 6 + 0] = r.get(0);
+            ret[i * 6 + 1] = r.get(1);
+            ret[i * 6 + 2] = r.get(2);
+            ret[i * 6 + 3] = t.get(0);
+            ret[i * 6 + 4] = t.get(1);
+            ret[i * 6 + 5] = t.get(2);
         }
 
-        return observations;
+        return ret;
     }
 
     /**
@@ -357,8 +368,7 @@
      * then calls {@link #mrcal_calibrate_camera} to perform calibration. Each corner's detection
      * level is converted to a weight (0.5^level), and negative levels indicate undetected corners.
      *
-     * @param observationCount Number of observations in `observationData`
-     * @param observationData An iterator of observations, each containing a list of corner locations,
+     * @param observations An list of observations, each containing a list of corner locations,
      *     decimation levels, and optional corner ids
      * @param boardWidth Number of internal corners horizontally
      * @param boardHeight Number of internal corners vertically
@@ -369,18 +379,110 @@
      * @return Calibration result with optimized intrinsics, poses, and error metrics
      */
     public static MrCalResult calibrateCamera(
-            int observationCount,
-            Iterator<MrCalObservation> observationData,
+            List<MrCalObservation> observations,
             int boardWidth,
             int boardHeight,
             double boardSpacing,
             int imageWidth,
             int imageHeight,
             double focalLen) {
 
-        var observations = makeObservations(observationCount, observationData, boardWidth, boardHeight);
+        var packedObservations = makeObservations(observations, boardWidth, boardHeight);
+
+        var results =
+                mrcal_calibrate_camera(
+                        packedObservations,
+                        boardWidth,
+                        boardHeight,
+                        boardSpacing,
+                        imageWidth,
+                        imageHeight,
+                        focalLen);
+
+        // Only return corners used for the corners that were provided in the input
+        for (int b = 0; b < observations.size(); b++) {
+            var observation = observations.get(b);
+            if (observation.ids != null) {
+                boolean[] fullCorners = results.cornersUsed.get(b);
+                boolean[] partialCorners = new boolean[observation.ids.length];
+                for (int i = 0; i < observation.ids.length; i++) {
+                    partialCorners[i] = fullCorners[observation.ids[i]];
+                }
+                results.cornersUsed.set(b, partialCorners);
+            }
+        }
+
+        return results;
+    }
+
+    /**
+     * High-level wrapper for computing uncertainty using mrcal.
+     *
+     * <p>Converts detected chessboard corners and their detection levels into a packed double array.
+     * Converts observation poses into a packed double array. Then calls {@link
+     * #mrcal_calibrate_camera} to perform calibration. Each corner's detection level is converted to
+     * a weight (0.5^level), and negative levels indicate undetected corners.
+     *
+     * @param observations An iterator of observations, each containing a list of corner locations,
+     *     decimation levels, and optional corner ids
+     * @param intrinsics Lens intrinsics from the calibration
+     * @param observationPoses Estimated pose of the calibration board in each snapshot
+     * @param boardWidth Number of internal corners horizontally
+     * @param boardHeight Number of internal corners vertically
+     * @param boardSpacing Physical spacing between corners (meters)
+     * @param imageWidth Image width in pixels
+     * @param imageHeight Image height in pixels
+     * @param sampleGridWidth Width of the generated sample grid
+     * @param sampleGridHeight Height of the generated sample grid
+     * @param warpX The x-component of lens distortion warp, or zero if no warp was estimated.
+     * @param warpY The y-component of lens distortion warp
+     * @return
+     */
+    public static Point3[] computeUncertainty(
+            List<MrCalObservation> observations,
+            double[] intrinsics,
+            List<Pose3d> observationPoses,
+            int boardWidth,
+            int boardHeight,
+            double boardSpacing,
+            int imageWidth,
+            int imageHeight,
+            int sampleGridWidth,
+            int sampleGridHeight,
+            double warpX,
+            double warpY) {
+        if (observations.size() != observationPoses.size()) {
+            // Mismatched sizes
+            return null;
+        }
+
+        var packedObservations = makeObservations(observations, boardWidth, boardHeight);
+
+        var packedResult =
+                compute_uncertainty(
+                        packedObservations,
+                        intrinsics,
+                        observationPosesToRtRef(observationPoses),
+                        boardWidth,
+                        boardHeight,
+                        boardSpacing,
+                        imageWidth,
+                        imageHeight,
+                        sampleGridWidth,
+                        sampleGridHeight,
+                        warpX,
+                        warpY);
+
+        var result = new Point3[sampleGridWidth * sampleGridHeight];
+
+        if (packedResult == null) {
+            return null;
+        }
+
+        for (int i = 0; i < result.length; i++) {
+            result[i] = new Point3(packedResult[i * 3], packedResult[i * 3 + 1], packedResult[i * 3 + 2]);
+        }
 
-        return mrcal_calibrate_camera(
-                observations, boardWidth, boardHeight, boardSpacing, imageWidth, imageHeight, focalLen);
+        return result;
     }
 }