Add withMultiTagFallbackStrategy()

Author: kcooney

vision/src/main/java/com/team2813/lib2813/vision/MultiPhotonPoseEstimator.java

@@ -67,14 +67,18 @@
 public class MultiPhotonPoseEstimator<C extends Camera> implements AutoCloseable {
   private final List<PhotonCameraWrapper<C>> cameraWrappers;
   private final Predicate<EstimatedRobotPose> isPoseValid;
+  private final PhotonPoseEstimator.PoseStrategy fallbackStrategy;
   private PhotonPoseEstimator.PoseStrategy poseEstimatorStrategy;
+  private boolean needHeadingData;
 
   /** A builder for {@code MultiPhotonPoseEstimator}. */
   public static final class Builder<C extends Camera> {
     private final Map<String, C> cameras = new HashMap<>();
     private final AprilTagFieldLayout aprilTagFieldLayout;
     private final NetworkTableInstance ntInstance;
     private final PhotonPoseEstimator.PoseStrategy poseEstimatorStrategy;
+    private PhotonPoseEstimator.PoseStrategy fallbackStrategy =
+        PhotonPoseEstimator.PoseStrategy.LOWEST_AMBIGUITY;
     private Predicate<EstimatedRobotPose> isPoseValid = pose -> true;
 
     Builder(
@@ -114,6 +118,24 @@ public Builder<C> withPoseFilter(Predicate<EstimatedRobotPose> isPoseValid) {
       return this;
     }
 
+    /**
+     * Set the Position Estimation Strategy used in multi-tag mode when only one tag can be seen.
+     * Must NOT be {@link PhotonPoseEstimator.PoseStrategy#MULTI_TAG_PNP_ON_COPROCESSOR} or {@link
+     * PhotonPoseEstimator.PoseStrategy#MULTI_TAG_PNP_ON_RIO}.
+     *
+     * <p>If this is not called, {@link PhotonPoseEstimator.PoseStrategy#LOWEST_AMBIGUITY} will be
+     * used.
+     *
+     * @param strategy the strategy to set
+     */
+    public Builder<C> withMultiTagFallbackStrategy(PhotonPoseEstimator.PoseStrategy strategy) {
+      if (isMultiTagStrategy(strategy)) {
+        throw new IllegalArgumentException("Fallback strategy cannot be " + strategy);
+      }
+      fallbackStrategy = strategy;
+      return this;
+    }
+
     /** Builds a configured MultiPhotonPoseEstimator. */
     public MultiPhotonPoseEstimator<C> build() {
       return new MultiPhotonPoseEstimator<>(this);
@@ -249,11 +271,13 @@ public void close() {
   /** Creates an instance using values from a {@code Builder}. */
   private MultiPhotonPoseEstimator(Builder<C> builder) {
     poseEstimatorStrategy = builder.poseEstimatorStrategy;
+    fallbackStrategy = builder.fallbackStrategy;
     isPoseValid = poseIsInField(builder.aprilTagFieldLayout).and(builder.isPoseValid);
     cameraWrappers =
         builder.cameras.values().stream()
             .map(camera -> createCameraWrapper(builder, camera))
             .collect(toCollection(ArrayList::new));
+    needHeadingData = calculateNeedHeadingData();
   }
 
   /**
@@ -268,6 +292,7 @@ private static <C extends Camera> PhotonCameraWrapper<C> createCameraWrapper(
     PhotonPoseEstimator estimator =
         new PhotonPoseEstimator(
             builder.aprilTagFieldLayout, builder.poseEstimatorStrategy, camera.robotToCamera());
+    estimator.setMultiTagFallbackStrategy(builder.fallbackStrategy);
 
     // Create NetworkTables publishers for 1) the position of the camera relative to the robot and
     // 2) the estimated position provided by the camera.
@@ -300,6 +325,7 @@ public void setPrimaryStrategy(PhotonPoseEstimator.PoseStrategy poseStrategy) {
     if (!poseStrategy.equals(poseEstimatorStrategy)) {
       cameraWrappers.forEach(wrapper -> wrapper.estimator.setPrimaryStrategy(poseStrategy));
       poseEstimatorStrategy = poseStrategy;
+      needHeadingData = calculateNeedHeadingData();
     }
   }
 
@@ -309,10 +335,7 @@ public void setPrimaryStrategy(PhotonPoseEstimator.PoseStrategy poseStrategy) {
    * @return {@code true} if the pose strategy is documented to require addHeadingData().
    */
   public boolean poseStrategyRequiresHeadingData() {
-    return switch (poseEstimatorStrategy) {
-      case PNP_DISTANCE_TRIG_SOLVE, CONSTRAINED_SOLVEPNP -> true;
-      default -> false;
-    };
+    return needHeadingData;
   }
 
   /**
@@ -428,6 +451,29 @@ public void close() {
     cameraWrappers.clear();
   }
 
+  private static boolean poseStrategyRequiresHeadingData(
+      PhotonPoseEstimator.PoseStrategy strategy) {
+    return switch (strategy) {
+      case PNP_DISTANCE_TRIG_SOLVE, CONSTRAINED_SOLVEPNP -> true;
+      default -> false;
+    };
+  }
+
+  private static boolean isMultiTagStrategy(PhotonPoseEstimator.PoseStrategy strategy) {
+    return switch (strategy) {
+      case MULTI_TAG_PNP_ON_COPROCESSOR, MULTI_TAG_PNP_ON_RIO -> true;
+      default -> false;
+    };
+  }
+
+  private boolean calculateNeedHeadingData() {
+    if (poseStrategyRequiresHeadingData(poseEstimatorStrategy)) {
+      return true;
+    }
+    return isMultiTagStrategy(poseEstimatorStrategy)
+        && poseStrategyRequiresHeadingData(fallbackStrategy);
+  }
+
   /** Creates a predicate that determines if a pose is inside the field. */
   private static <C extends Camera> Predicate<EstimatedRobotPose> poseIsInField(
       AprilTagFieldLayout aprilTagFieldLayout) {