Return partial used corners for calibration

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

@@ -19,7 +19,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Iterator;
 import java.util.List;
 import org.opencv.core.Point;
 import org.wpilib.math.geometry.Pose3d;
@@ -27,12 +26,12 @@
 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,16 +56,16 @@
         }
     }
 
     public static class MrCalResult {
         public boolean success;
         public double[] intrinsics;
         public double rms_error;
         public double[] residuals;
         public double warp_x;
         public double warp_y;
         public int Noutliers;
         public List<Pose3d> optimizedPoses;
         public List<boolean[]> cornersUsed;
 
         public MrCalResult(boolean success) {
             this.success = success;
@@ -272,27 +271,20 @@
             double warpY);
 
     /**
-     * Convert an iterator of board-pixel-corners, detection decimation levels, and corner IDs for
-     * each snapshot of a chessboard boardWidth x boardHeight to a packed double[] suitable to pass to
-     * MrCalJni::mrcal_calibrate_camera. Ids will be used to select which corners are actually
-     * present. Levels will be converted to weights using weight = 0.5^level, as explained
+     * Convert a list of board-pixel-corners, detection decimation levels, and corner IDs for each
+     * snapshot of a chessboard boardWidth x boardHeight to a packed double[] suitable to pass to
+     * MrCalJni::mrcal_calibrate_camera. Ids are used to determine which corners are actually present
+     * and so the returned usage info matches the caller's input. Levels will be converted to weights
+     * using weight = 0.5^level, as explained
      * [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,14 @@
 
                     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 observations;
+        return packedObservations;
     }
 
     /**
@@ -357,9 +344,13 @@
      * 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,
-     *     decimation levels, and optional corner ids
+     * <p>When observations include corner IDs, the returned corner-usage mask is remapped back to the
+     * same subset and order that was provided by the caller.
+     *
+     * @param observations An list of observations, each containing a list of corner locations,
+     *     decimation levels, and optional corner ids. If ids is null, the observation is treated as a
+     *     full board. If ids is present, only the listed corners are used; each id must be within the
+     *     board bounds and non-negative.
      * @param boardWidth Number of internal corners horizontally
      * @param boardHeight Number of internal corners vertically
      * @param boardSpacing Physical spacing between corners (meters)
@@ -369,18 +360,39 @@
      * @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 mrcal_calibrate_camera(
-                observations, boardWidth, boardHeight, boardSpacing, imageWidth, imageHeight, focalLen);
+        return results;
     }
 }