Source code

001/************************ PROJECT ZIXI ************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved.*/
003/* This work is licensed under the terms of the MIT license.  */
004/**************************************************************/
005package com.stuypulse.robot.subsystems.vision;
006
007import com.stuypulse.robot.constants.FieldConstants;
008import com.stuypulse.robot.constants.GlobalSettings;
009import com.stuypulse.robot.subsystems.vision.VisionConstants.Camera;
010import com.stuypulse.robot.subsystems.vision.VisionConstants.CameraData;
011import com.stuypulse.robot.subsystems.vision.VisionConstants.Pipeline;
012import com.stuypulse.robot.subsystems.vision.VisionConstants.VisionSettings;
013import com.stuypulse.robot.subsystems.vision.VisionIO.MegaTagMode;
014import com.stuypulse.robot.subsystems.vision.VisionIO.PoseObservation;
015import com.stuypulse.robot.subsystems.vision.VisionIO.PoseObservationType;
016import com.stuypulse.robot.subsystems.vision.VisionIO.VisionIOOutputs;
017import com.stuypulse.robot.util.PeriodicSubsystem;
018
019import org.wpilib.command3.Command;
020import org.wpilib.math.geometry.Pose2d;
021import org.wpilib.math.geometry.Pose3d;
022import org.wpilib.math.geometry.Rotation2d;
023import org.wpilib.math.linalg.Matrix;
024import org.wpilib.math.linalg.VecBuilder;
025import org.wpilib.math.numbers.N1;
026import org.wpilib.math.numbers.N3;
027import org.wpilib.system.Timer;
028import org.wpilib.util.Alert;
029
030import java.util.EnumMap;
031import java.util.Iterator;
032import java.util.LinkedList;
033import java.util.List;
034import java.util.Map.Entry;
035import java.util.Optional;
036import org.littletonrobotics.junction.Logger;
037
038public class Vision extends PeriodicSubsystem {
039    private final VisionConsumer consumer;
040    private final EnumMap<Camera, VisionIO> io;
041    private final EnumMap<Camera, VisionIOInputsAutoLogged> inputs;
042    private final EnumMap<Camera, VisionIOOutputs> outputs;
043    private final EnumMap<Camera, Alert> disconnectedAlerts;
044
045    private final Timer hdrTimer;
046
047    public Vision(VisionConsumer consumer, EnumMap<Camera, VisionIO> cameraIOMap) {
048        this.consumer = consumer;
049
050        this.io = new EnumMap<>(cameraIOMap);
051
052        // Initialize inputs
053        this.inputs = new EnumMap<>(Camera.class);
054        this.outputs = new EnumMap<>(Camera.class);
055
056        // Initialize disconnected alerts
057        this.disconnectedAlerts = new EnumMap<>(Camera.class);
058
059        this.hdrTimer = new Timer();
060        for (Camera camera : io.keySet()) {
061            inputs.put(camera, new VisionIOInputsAutoLogged());
062            outputs.put(camera, new VisionIOOutputs());
063            disconnectedAlerts.put(
064                    camera,
065                    new Alert(
066                            "Vision " + camera.getName(),
067                            "Vision camera " + camera.getName() + " is disconnected.",
068                            Alert.Level.MEDIUM));
069        }
070
071        hdrTimer.start();
072    }
073
074    /**
075     * Returns the X angle to the best target, which can be used for simple servoing with vision.
076     *
077     * @param camera The camera to use
078     */
079    public Rotation2d getTargetX(Camera camera) {
080        return inputs.get(camera).latestTargetObservation.tx();
081    }
082
083    public Optional<Integer> getPipeline() {
084        Iterator<VisionIOInputsAutoLogged> iterator = inputs.values().iterator();
085        return iterator.hasNext() ? Optional.of(iterator.next().currentPipeline) : Optional.empty();
086    }
087
088    @Override
089    public void periodic() {
090        if (!GlobalSettings.EnabledSubsystems.VISION.get()) {
091            return;
092        }
093
094        for (Entry<Camera, VisionIO> entry : io.entrySet()) {
095            VisionIO currentIO = entry.getValue();
096            VisionIOInputsAutoLogged currentInputs = inputs.get(entry.getKey());
097            currentIO.updateInputs(currentInputs);
098            Logger.processInputs("Vision/" + entry.getKey().getName(), currentInputs);
099        }
100
101        Optional<Integer> pipeline = getPipeline();
102
103        if (hdrTimer.hasElapsed(VisionConstants.VisionSettings.HDR_SWITCHING_INTERVAL)
104                && VisionConstants.VisionSettings.HDR_ENABLED.get()
105                && pipeline.isPresent()) {
106            int nextPipeline =
107                    pipeline.get() == Pipeline.HIGH_SUN.getIndex()
108                            ? Pipeline.LOW_SUN.getIndex()
109                            : Pipeline.HIGH_SUN.getIndex();
110
111            for (VisionIOOutputs outputs : outputs.values()) {
112                outputs.pipeline = nextPipeline;
113            }
114
115            hdrTimer.reset();
116        }
117
118        // Initialize logging values
119        List<Pose3d> allTagPoses = new LinkedList<>();
120        List<Pose3d> allRobotPoses = new LinkedList<>();
121        List<Pose3d> allRobotPosesAccepted = new LinkedList<>();
122        List<Pose3d> allRobotPosesRejected = new LinkedList<>();
123
124        // Loop over cameras
125        for (Entry<Camera, VisionIO> entry : io.entrySet()) {
126            CameraData currentCameraData = entry.getKey().getData();
127            VisionIOInputsAutoLogged currentInputs = inputs.get(entry.getKey());
128
129            // Update disconnected alert
130            disconnectedAlerts.get(entry.getKey()).set(!currentInputs.connected);
131
132            // Initialize logging values
133            List<Pose3d> tagPoses = new LinkedList<>();
134            List<Pose3d> robotPoses = new LinkedList<>();
135            List<Pose3d> robotPosesAccepted = new LinkedList<>();
136            List<Pose3d> robotPosesRejected = new LinkedList<>();
137
138            // Add tag poses
139            for (int tagId : currentInputs.tagIds) {
140                Optional<Pose3d> tagPose = FieldConstants.FIELD_LAYOUT.getTagPose(tagId);
141                if (tagPose.isPresent()) {
142                    tagPoses.add(tagPose.get());
143                }
144            }
145
146            // Loop over pose observations
147            for (PoseObservation observation : currentInputs.poseObservations) {
148                // Check whether to reject pose
149                boolean rejectPose =
150                        observation.tagCount() == 0 // Must have at least one tag
151                                || (observation.tagCount() == 1
152                                        && observation.ambiguity()
153                                                > VisionSettings
154                                                        .MAX_AMBIGUITY) // Cannot be high ambiguity
155                                || Math.abs(observation.pose().getZ())
156                                        > VisionSettings
157                                                .MAX_Z_ERROR // Must have realistic Z coordinate
158
159                                // Must be within the field boundaries
160                                || observation.pose().getX() < 0.0
161                                || observation.pose().getX()
162                                        > FieldConstants.FIELD_LAYOUT.getFieldLength()
163                                || observation.pose().getY() < 0.0
164                                || observation.pose().getY()
165                                        > FieldConstants.FIELD_LAYOUT.getFieldWidth();
166
167                // Add pose to log
168                robotPoses.add(observation.pose());
169                if (rejectPose) {
170                    robotPosesRejected.add(observation.pose());
171                } else {
172                    robotPosesAccepted.add(observation.pose());
173                }
174
175                // Skip if rejected
176                if (rejectPose) {
177                    continue;
178                }
179
180                // Calculate standard deviations
181                double stdDevFactor =
182                        Math.pow(observation.averageTagDistance(), 2.0) / observation.tagCount();
183                double linearStdDev = VisionSettings.LINEAR_STD_DEV_BASELINE * stdDevFactor;
184                double angularStdDev = VisionSettings.ANGULAR_STD_DEV_BASELINE * stdDevFactor;
185                if (observation.type() == PoseObservationType.MEGATAG_2) {
186                    linearStdDev *= VisionSettings.LINEAR_STD_DEV_MEGATAG_2_FACTOR;
187                    angularStdDev *= VisionSettings.ANGULAR_STD_DEV_MEGATAG_2_FACTOR;
188                }
189                linearStdDev *= currentCameraData.stdDevFactor();
190                angularStdDev *= currentCameraData.stdDevFactor();
191
192                // Send vision observation
193                consumer.accept(
194                        observation.pose().toPose2d(),
195                        observation.timestamp(),
196                        VecBuilder.fill(linearStdDev, linearStdDev, angularStdDev));
197            }
198
199            // Log camera metadata
200            Logger.recordOutput(
201                    "Vision/Camera" + currentCameraData.name() + "/TagPoses",
202                    tagPoses.toArray(new Pose3d[tagPoses.size()]));
203            Logger.recordOutput(
204                    "Vision/Camera" + currentCameraData.name() + "/RobotPoses",
205                    robotPoses.toArray(new Pose3d[robotPoses.size()]));
206            Logger.recordOutput(
207                    "Vision/Camera" + currentCameraData.name() + "/RobotPosesAccepted",
208                    robotPosesAccepted.toArray(new Pose3d[robotPosesAccepted.size()]));
209            Logger.recordOutput(
210                    "Vision/Camera" + currentCameraData.name() + "/RobotPosesRejected",
211                    robotPosesRejected.toArray(new Pose3d[robotPosesRejected.size()]));
212            allTagPoses.addAll(tagPoses);
213            allRobotPoses.addAll(robotPoses);
214            allRobotPosesAccepted.addAll(robotPosesAccepted);
215            allRobotPosesRejected.addAll(robotPosesRejected);
216        }
217
218        // Log summary data
219        Logger.recordOutput("Vision/Summary/TagPoses", allTagPoses.toArray(new Pose3d[0]));
220        Logger.recordOutput("Vision/Summary/RobotPoses", allRobotPoses.toArray(new Pose3d[0]));
221        Logger.recordOutput(
222                "Vision/Summary/RobotPosesAccepted", allRobotPosesAccepted.toArray(new Pose3d[0]));
223        Logger.recordOutput(
224                "Vision/Summary/RobotPosesRejected", allRobotPosesRejected.toArray(new Pose3d[0]));
225    }
226
227    @Override
228    public void periodicAfterScheduler() {
229        for (Entry<Camera, VisionIO> entry : io.entrySet()) {
230            entry.getValue().applyOutputs(outputs.get(entry.getKey()));
231        }
232    }
233
234    @FunctionalInterface
235    public static interface VisionConsumer {
236        public void accept(
237                Pose2d visionRobotPoseMeters,
238                double timestampSeconds,
239                Matrix<N3, N1> visionMeasurementStdDevs);
240    }
241
242    public void enable() {
243        GlobalSettings.EnabledSubsystems.VISION.set(true);
244    }
245
246    public void disable() {
247        GlobalSettings.EnabledSubsystems.VISION.set(false);
248    }
249
250    public Command setPipeline(int pipeline) {
251        return run(coroutine -> {
252                    for (VisionIOOutputs output : outputs.values()) {
253                        output.pipeline = pipeline;
254                    }
255                })
256                .named("Vision Set Pipeline");
257    }
258
259    public Command setMegaTagMode(MegaTagMode mode) {
260        return run(coroutine -> {
261                    for (VisionIOOutputs output : outputs.values()) {
262                        output.megaTagMode = mode;
263                    }
264                })
265                .named("Vision Set MegaTag Mode");
266    }
267
268    public Command setVisionEnabled() {
269        return run(coroutine -> {
270                    this.enable();
271                })
272                .named("Vision Set Enabled");
273    }
274
275    public Command setVisionDisabled() {
276        return run(coroutine -> {
277                    this.disable();
278                })
279                .named("Vision Set Enabled");
280    }
281
282    public Command setAprilTagWhitelist(int[] whitelist) {
283        return run(coroutine -> {
284                    for (VisionIOOutputs output : outputs.values()) {
285                        output.aprilTagWhitelist = whitelist.clone();
286                    }
287                })
288                .named("Vision Set April Tag Whitelist");
289    }
290
291    private Command setIMUMode(int imuMode, String name) {
292        return run(coroutine -> {
293                    for (VisionIOOutputs output : outputs.values()) {
294                        output.imuMode = imuMode;
295                    }
296                })
297                .named(name);
298    }
299
300    public Command setIMUMode(int imuMode) {
301        return setIMUMode(imuMode, "Vision Set IMU Mode");
302    }
303
304    public Command resetIMU() {
305        return setIMUMode(VisionSettings.RESET_IMU_INDEX, "Vision Reset IMU");
306    }
307}