From 975519ac1012ba874901ff869acc43b029cf262c Mon Sep 17 00:00:00 2001
From: amquake <noleetarrr@gmail.com>
Date: Thu, 19 Oct 2023 23:06:36 -0700
Subject: [PATCH] multitag support

---
 .../components/dashboard/tabs/ArucoTab.vue    |  10 +-
 .../components/dashboard/tabs/OutputTab.vue   |  10 +-
 .../components/dashboard/tabs/TargetsTab.vue  |   3 +-
 photon-client/src/types/PipelineTypes.ts      |  10 +-
 .../pipeline/AprilTagPipelineSettings.java    |   2 +-
 .../vision/pipeline/ArucoPipeline.java        | 122 +++++++++++++-----
 .../pipeline/ArucoPipelineSettings.java       |   3 +
 7 files changed, 121 insertions(+), 39 deletions(-)

diff --git a/photon-client/src/components/dashboard/tabs/ArucoTab.vue b/photon-client/src/components/dashboard/tabs/ArucoTab.vue
index c0340fa99f..3bdbd5d8ce 100644
--- a/photon-client/src/components/dashboard/tabs/ArucoTab.vue
+++ b/photon-client/src/components/dashboard/tabs/ArucoTab.vue
@@ -4,6 +4,7 @@ import { PipelineType } from "@/types/PipelineTypes";
 import PvSlider from "@/components/common/pv-slider.vue";
 import pvSwitch from "@/components/common/pv-switch.vue";
 import pvRangeSlider from "@/components/common/pv-range-slider.vue";
+import pvSelect from "@/components/common/pv-select.vue";
 import { computed, getCurrentInstance } from "vue";
 import { useStateStore } from "@/stores/StateStore";
 
@@ -11,7 +12,7 @@ import { useStateStore } from "@/stores/StateStore";
 // Defer reference to store access method
 const currentPipelineSettings = useCameraSettingsStore().currentPipelineSettings;
 
-// TODO fix cv-range-slider so that store access doesn't need to be deferred
+// TODO fix pv-range-slider so that store access doesn't need to be deferred
 const threshWinSizes = computed<[number, number]>({
   get: () => {
     if (currentPipelineSettings.pipelineType === PipelineType.Aruco) {
@@ -38,6 +39,13 @@ const interactiveCols = computed(
 
 <template>
   <div v-if="currentPipelineSettings.pipelineType === PipelineType.Aruco">
+    <pv-select
+      v-model="currentPipelineSettings.tagFamily"
+      label="Target family"
+      :items="['AprilTag Family 36h11', 'AprilTag Family 25h9', 'AprilTag Family 16h5']"
+      :select-cols="interactiveCols"
+      @input="(value) => useCameraSettingsStore().changeCurrentPipelineSetting({ tagFamily: value }, false)"
+    />
     <pv-switch
       v-model="currentPipelineSettings.useCornerRefinement"
       class="pt-2"
diff --git a/photon-client/src/components/dashboard/tabs/OutputTab.vue b/photon-client/src/components/dashboard/tabs/OutputTab.vue
index 5116b6ee38..b8b0e5f340 100644
--- a/photon-client/src/components/dashboard/tabs/OutputTab.vue
+++ b/photon-client/src/components/dashboard/tabs/OutputTab.vue
@@ -84,9 +84,10 @@ const interactiveCols = computed(
         (value) => useCameraSettingsStore().changeCurrentPipelineSetting({ outputShowMultipleTargets: value }, false)
       "
     />
-    <cv-switch
+    <pv-switch
       v-if="
-        currentPipelineSettings.pipelineType === PipelineType.AprilTag &&
+        (currentPipelineSettings.pipelineType === PipelineType.AprilTag ||
+        currentPipelineSettings.pipelineType === PipelineType.Aruco) &&
         useCameraSettingsStore().isCurrentVideoFormatCalibrated
       "
       v-model="currentPipelineSettings.doMultiTarget"
@@ -96,9 +97,10 @@ const interactiveCols = computed(
       :disabled="!isTagPipeline"
       @input="(value) => useCameraSettingsStore().changeCurrentPipelineSetting({ doMultiTarget: value }, false)"
     />
-    <cv-switch
+    <pv-switch
       v-if="
-        currentPipelineSettings.pipelineType === PipelineType.AprilTag &&
+        (currentPipelineSettings.pipelineType === PipelineType.AprilTag ||
+        currentPipelineSettings.pipelineType === PipelineType.Aruco) &&
         useCameraSettingsStore().isCurrentVideoFormatCalibrated
       "
       v-model="currentPipelineSettings.doSingleTargetAlways"
diff --git a/photon-client/src/components/dashboard/tabs/TargetsTab.vue b/photon-client/src/components/dashboard/tabs/TargetsTab.vue
index d0e31d5f13..c5cc5a8fde 100644
--- a/photon-client/src/components/dashboard/tabs/TargetsTab.vue
+++ b/photon-client/src/components/dashboard/tabs/TargetsTab.vue
@@ -80,7 +80,8 @@ import { useStateStore } from "@/stores/StateStore";
     </v-row>
     <v-row
       v-if="
-        useCameraSettingsStore().currentPipelineSettings.pipelineType === PipelineType.AprilTag &&
+        (useCameraSettingsStore().currentPipelineType === PipelineType.AprilTag ||
+                    useCameraSettingsStore().currentPipelineType === PipelineType.Aruco) &&
         useCameraSettingsStore().currentPipelineSettings.doMultiTarget
       "
       align="start"
diff --git a/photon-client/src/types/PipelineTypes.ts b/photon-client/src/types/PipelineTypes.ts
index 19dcd8dad1..d6debea90d 100644
--- a/photon-client/src/types/PipelineTypes.ts
+++ b/photon-client/src/types/PipelineTypes.ts
@@ -246,6 +246,8 @@ export const DefaultAprilTagPipelineSettings: AprilTagPipelineSettings = {
 export interface ArucoPipelineSettings extends PipelineSettings {
   pipelineType: PipelineType.Aruco;
 
+  tagFamily: AprilTagFamily;
+
   threshWinSizes: WebsocketNumberPair | [number, number];
   threshStepSize: number;
   threshConstant: number;
@@ -256,6 +258,9 @@ export interface ArucoPipelineSettings extends PipelineSettings {
   useAruco3: boolean;
   aruco3MinMarkerSideRatio: number;
   aruco3MinCanonicalImgSide: number;
+
+  doMultiTarget: boolean;
+  doSingleTargetAlways: boolean;
 }
 export type ConfigurableArucoPipelineSettings = Partial<Omit<ArucoPipelineSettings, "pipelineType">> &
   ConfigurablePipelineSettings;
@@ -268,6 +273,7 @@ export const DefaultArucoPipelineSettings: ArucoPipelineSettings = {
   ledMode: false,
   pipelineType: PipelineType.Aruco,
 
+  tagFamily: AprilTagFamily.Family16h5,
   threshWinSizes: { first: 11, second: 91 },
   threshStepSize: 40,
   threshConstant: 10,
@@ -275,7 +281,9 @@ export const DefaultArucoPipelineSettings: ArucoPipelineSettings = {
   useCornerRefinement: true,
   useAruco3: false,
   aruco3MinMarkerSideRatio: 0.02,
-  aruco3MinCanonicalImgSide: 32
+  aruco3MinCanonicalImgSide: 32,
+  doMultiTarget: false,
+  doSingleTargetAlways: false
 };
 
 export type ActivePipelineSettings =
diff --git a/photon-core/src/main/java/org/photonvision/vision/pipeline/AprilTagPipelineSettings.java b/photon-core/src/main/java/org/photonvision/vision/pipeline/AprilTagPipelineSettings.java
index 47faf82e29..6db514d005 100644
--- a/photon-core/src/main/java/org/photonvision/vision/pipeline/AprilTagPipelineSettings.java
+++ b/photon-core/src/main/java/org/photonvision/vision/pipeline/AprilTagPipelineSettings.java
@@ -32,7 +32,7 @@ public class AprilTagPipelineSettings extends AdvancedPipelineSettings {
     public int numIterations = 40;
     public int hammingDist = 0;
     public int decisionMargin = 35;
-    public boolean doMultiTarget = true;
+    public boolean doMultiTarget = false;
     public boolean doSingleTargetAlways = false;
 
     // 3d settings
diff --git a/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipeline.java b/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipeline.java
index 41e3e404d3..19f92c0e85 100644
--- a/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipeline.java
+++ b/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipeline.java
@@ -35,6 +35,8 @@
 package org.photonvision.vision.pipeline;
 
 import edu.wpi.first.apriltag.AprilTagPoseEstimate;
+import edu.wpi.first.math.geometry.CoordinateSystem;
+import edu.wpi.first.math.geometry.Pose3d;
 import edu.wpi.first.math.geometry.Transform3d;
 import edu.wpi.first.math.util.Units;
 import java.util.ArrayList;
@@ -42,21 +44,24 @@
 import org.opencv.core.Mat;
 import org.opencv.imgproc.Imgproc;
 import org.opencv.objdetect.Objdetect;
+import org.photonvision.common.configuration.ConfigManager;
 import org.photonvision.common.util.math.MathUtils;
-import org.photonvision.vision.apriltag.AprilTagFamily;
+import org.photonvision.targeting.MultiTargetPNPResults;
 import org.photonvision.vision.aruco.ArucoDetectionResult;
 import org.photonvision.vision.frame.Frame;
 import org.photonvision.vision.frame.FrameThresholdType;
 import org.photonvision.vision.pipe.CVPipe.CVPipeResult;
 import org.photonvision.vision.pipe.impl.*;
 import org.photonvision.vision.pipe.impl.ArucoPoseEstimatorPipe.ArucoPoseEstimatorPipeParams;
+import org.photonvision.vision.pipe.impl.MultiTargetPNPPipe.MultiTargetPNPPipeParams;
 import org.photonvision.vision.pipeline.result.CVPipelineResult;
 import org.photonvision.vision.target.TrackedTarget;
 import org.photonvision.vision.target.TrackedTarget.TargetCalculationParameters;
 
 public class ArucoPipeline extends CVPipeline<CVPipelineResult, ArucoPipelineSettings> {
     private final ArucoDetectionPipe arucoDetectionPipe = new ArucoDetectionPipe();
-    private final ArucoPoseEstimatorPipe poseEstimatorPipe = new ArucoPoseEstimatorPipe();
+    private final ArucoPoseEstimatorPipe singleTagPoseEstimatorPipe = new ArucoPoseEstimatorPipe();
+    private final MultiTargetPNPPipe multiTagPNPPipe = new MultiTargetPNPPipe();
     private final CalculateFPSPipe calculateFPSPipe = new CalculateFPSPipe();
 
     public ArucoPipeline() {
@@ -74,8 +79,15 @@ protected void setPipeParamsImpl() {
         var params = new ArucoDetectionPipeParams();
         // sanitize and record settings
 
-        if (settings.tagFamily == AprilTagFamily.kTag36h11) {
-            params.tagFamily = Objdetect.DICT_APRILTAG_36h11;
+        switch(settings.tagFamily) {
+            case kTag36h11:
+                params.tagFamily = Objdetect.DICT_APRILTAG_36h11;
+                break;
+            case kTag25h9:
+                params.tagFamily = Objdetect.DICT_APRILTAG_25h9;
+                break;
+            default:
+                params.tagFamily = Objdetect.DICT_APRILTAG_16h5;
         }
 
         int threshMinSize = Math.max(3, settings.threshWinSizes.getFirst());
@@ -99,11 +111,16 @@ protected void setPipeParamsImpl() {
 
         if (frameStaticProperties.cameraCalibration != null) {
             var cameraMatrix = frameStaticProperties.cameraCalibration.getCameraIntrinsicsMat();
-            if (cameraMatrix != null) {
+            if (cameraMatrix != null && cameraMatrix.rows() > 0) {
                 var estimatorParams =
                         new ArucoPoseEstimatorPipeParams(
                                 frameStaticProperties.cameraCalibration, Units.inchesToMeters(6));
-                poseEstimatorPipe.setParams(estimatorParams);
+                singleTagPoseEstimatorPipe.setParams(estimatorParams);
+
+                // TODO global state ew
+                var atfl = ConfigManager.getInstance().getConfig().getApriltagFieldLayout();
+                multiTagPNPPipe.setParams(
+                        new MultiTargetPNPPipeParams(frameStaticProperties.cameraCalibration, atfl));
             }
         }
     }
@@ -112,8 +129,6 @@ protected void setPipeParamsImpl() {
     protected CVPipelineResult process(Frame frame, ArucoPipelineSettings settings) {
         long sumPipeNanosElapsed = 0L;
 
-        List<TrackedTarget> targetList;
-
         if (frame.type != FrameThresholdType.GREYSCALE) {
             // We asked for a GREYSCALE frame, but didn't get one -- best we can do is give up
             return new CVPipelineResult(0, 0, List.of(), frame);
@@ -132,43 +147,88 @@ protected CVPipelineResult process(Frame frame, ArucoPipelineSettings settings)
                     settings.threshConstant);
         }
 
-        targetList = new ArrayList<>();
+        List<TrackedTarget> targetList = new ArrayList<>();
         for (ArucoDetectionResult detection : tagDetectionPipeResult.output) {
-            // TODO this should be in a pipe, not in the top level here (Matt)
-
-            AprilTagPoseEstimate tagPoseEstimate = null;
-            if (settings.solvePNPEnabled) {
-                var poseResult = poseEstimatorPipe.run(detection);
-                sumPipeNanosElapsed += poseResult.nanosElapsed;
-                tagPoseEstimate = poseResult.output;
-            }
-
-            // populate the target list
-            // Challenge here is that TrackedTarget functions with OpenCV Contour
+            // Populate target list for multitag
+            // (TODO: Address circular dependencies. Multitag only requires corners and IDs, this should
+            // not be necessary.)
             TrackedTarget target =
                     new TrackedTarget(
                             detection,
-                            tagPoseEstimate,
+                            null,
                             new TargetCalculationParameters(
                                     false, null, null, null, null, frameStaticProperties));
 
-            var correctedBestPose =
-                    MathUtils.convertOpenCVtoPhotonTransform(target.getBestCameraToTarget3d());
-            var correctedAltPose =
-                    MathUtils.convertOpenCVtoPhotonTransform(target.getAltCameraToTarget3d());
+            targetList.add(target);
+        }
 
-            target.setBestCameraToTarget3d(
-                    new Transform3d(correctedBestPose.getTranslation(), correctedBestPose.getRotation()));
-            target.setAltCameraToTarget3d(
-                    new Transform3d(correctedAltPose.getTranslation(), correctedAltPose.getRotation()));
+        // Do multi-tag pose estimation
+        MultiTargetPNPResults multiTagResult = new MultiTargetPNPResults();
+        if (settings.solvePNPEnabled && settings.doMultiTarget) {
+            var multiTagOutput = multiTagPNPPipe.run(targetList);
+            sumPipeNanosElapsed += multiTagOutput.nanosElapsed;
+            multiTagResult = multiTagOutput.output;
+        }
 
-            targetList.add(target);
+        // Do single-tag pose estimation
+        if (settings.solvePNPEnabled) {
+            // Clear target list that was used for multitag so we can add target transforms
+            targetList.clear();
+            // TODO global state again ew
+            var atfl = ConfigManager.getInstance().getConfig().getApriltagFieldLayout();
+
+            for (ArucoDetectionResult detection : tagDetectionPipeResult.output) {
+                AprilTagPoseEstimate tagPoseEstimate = null;
+                // Do single-tag estimation when "always enabled" or if a tag was not used for multitag
+                if (settings.doSingleTargetAlways
+                        || !multiTagResult.fiducialIDsUsed.contains(Integer.valueOf(detection.getId()))) {
+                    var poseResult = singleTagPoseEstimatorPipe.run(detection);
+                    sumPipeNanosElapsed += poseResult.nanosElapsed;
+                    tagPoseEstimate = poseResult.output;
+                }
+
+                // If single-tag estimation was not done, this is a multi-target tag from the layout
+                if (tagPoseEstimate == null) {
+                    // compute this tag's camera-to-tag transform using the multitag result
+                    var tagPose = atfl.getTagPose(detection.getId());
+                    if (tagPose.isPresent()) {
+                        var camToTag =
+                                new Transform3d(
+                                        new Pose3d().plus(multiTagResult.estimatedPose.best), tagPose.get());
+                        // match expected OpenCV coordinate system
+                        camToTag = CoordinateSystem.convert(camToTag, CoordinateSystem.NWU(), CoordinateSystem.EDN());
+
+                        tagPoseEstimate = new AprilTagPoseEstimate(camToTag, camToTag, 0, 0);
+                    }
+                }
+
+                // populate the target list
+                // Challenge here is that TrackedTarget functions with OpenCV Contour
+                TrackedTarget target =
+                        new TrackedTarget(
+                                detection,
+                                tagPoseEstimate,
+                                new TargetCalculationParameters(
+                                        false, null, null, null, null, frameStaticProperties));
+
+                var correctedBestPose =
+                        MathUtils.convertOpenCVtoPhotonTransform(target.getBestCameraToTarget3d());
+                var correctedAltPose =
+                        MathUtils.convertOpenCVtoPhotonTransform(target.getAltCameraToTarget3d());
+
+                target.setBestCameraToTarget3d(
+                        new Transform3d(correctedBestPose.getTranslation(), correctedBestPose.getRotation()));
+                target.setAltCameraToTarget3d(
+                        new Transform3d(correctedAltPose.getTranslation(), correctedAltPose.getRotation()));
+
+                targetList.add(target);
+            }
         }
 
         var fpsResult = calculateFPSPipe.run(null);
         var fps = fpsResult.output;
 
-        return new CVPipelineResult(sumPipeNanosElapsed, fps, targetList, frame);
+        return new CVPipelineResult(sumPipeNanosElapsed, fps, targetList, multiTagResult, frame);
     }
 
     private void drawThresholdFrame(Mat greyMat, Mat outputMat, int windowSize, double constant) {
diff --git a/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipelineSettings.java b/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipelineSettings.java
index 61d398621a..94eb963435 100644
--- a/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipelineSettings.java
+++ b/photon-core/src/main/java/org/photonvision/vision/pipeline/ArucoPipelineSettings.java
@@ -39,6 +39,9 @@ public class ArucoPipelineSettings extends AdvancedPipelineSettings {
     public double aruco3MinMarkerSideRatio = 0.02;
     public int aruco3MinCanonicalImgSide = 32;
 
+    public boolean doMultiTarget = false;
+    public boolean doSingleTargetAlways = false;
+
     public ArucoPipelineSettings() {
         super();
         pipelineType = PipelineType.Aruco;