001/*
002 * Copyright (C) Photon Vision.
003 *
004 * This program is free software: you can redistribute it and/or modify
005 * it under the terms of the GNU General Public License as published by
006 * the Free Software Foundation, either version 3 of the License, or
007 * (at your option) any later version.
008 *
009 * This program is distributed in the hope that it will be useful,
010 * but WITHOUT ANY WARRANTY; without even the implied warranty of
011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
012 * GNU General Public License for more details.
013 *
014 * You should have received a copy of the GNU General Public License
015 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
016 */
017
018/*
019 * Copyright (C) Photon Vision.
020 *
021 * This program is free software: you can redistribute it and/or modify
022 * it under the terms of the GNU General Public License as published by
023 * the Free Software Foundation, either version 3 of the License, or
024 * (at your option) any later version.
025 *
026 * This program is distributed in the hope that it will be useful,
027 * but WITHOUT ANY WARRANTY; without even the implied warranty of
028 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
029 * GNU General Public License for more details.
030 *
031 * You should have received a copy of the GNU General Public License
032 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
033 */
034
035package org.photonvision.vision.pipeline;
036
037import edu.wpi.first.apriltag.AprilTagPoseEstimate;
038import edu.wpi.first.math.geometry.CoordinateSystem;
039import edu.wpi.first.math.geometry.Pose3d;
040import edu.wpi.first.math.geometry.Transform3d;
041import edu.wpi.first.math.util.Units;
042import java.util.ArrayList;
043import java.util.List;
044import java.util.Optional;
045import org.opencv.core.Mat;
046import org.opencv.imgproc.Imgproc;
047import org.opencv.objdetect.Objdetect;
048import org.photonvision.common.configuration.ConfigManager;
049import org.photonvision.common.util.math.MathUtils;
050import org.photonvision.estimation.TargetModel;
051import org.photonvision.targeting.MultiTargetPNPResult;
052import org.photonvision.vision.aruco.ArucoDetectionResult;
053import org.photonvision.vision.frame.Frame;
054import org.photonvision.vision.frame.FrameThresholdType;
055import org.photonvision.vision.pipe.CVPipe.CVPipeResult;
056import org.photonvision.vision.pipe.impl.*;
057import org.photonvision.vision.pipe.impl.ArucoPoseEstimatorPipe.ArucoPoseEstimatorPipeParams;
058import org.photonvision.vision.pipe.impl.MultiTargetPNPPipe.MultiTargetPNPPipeParams;
059import org.photonvision.vision.pipeline.result.CVPipelineResult;
060import org.photonvision.vision.target.TrackedTarget;
061import org.photonvision.vision.target.TrackedTarget.TargetCalculationParameters;
062
063public class ArucoPipeline extends CVPipeline<CVPipelineResult, ArucoPipelineSettings> {
064    private ArucoDetectionPipe arucoDetectionPipe = new ArucoDetectionPipe();
065    private ArucoPoseEstimatorPipe singleTagPoseEstimatorPipe = new ArucoPoseEstimatorPipe();
066    private final MultiTargetPNPPipe multiTagPNPPipe = new MultiTargetPNPPipe();
067    private final CalculateFPSPipe calculateFPSPipe = new CalculateFPSPipe();
068
069    public ArucoPipeline() {
070        super(FrameThresholdType.GREYSCALE);
071        settings = new ArucoPipelineSettings();
072    }
073
074    public ArucoPipeline(ArucoPipelineSettings settings) {
075        super(FrameThresholdType.GREYSCALE);
076        this.settings = settings;
077    }
078
079    @Override
080    protected void setPipeParamsImpl() {
081        var params = new ArucoDetectionPipeParams();
082        // sanitize and record settings
083
084        // for now, hard code tag width based on enum value
085        // 2023/other: best guess is 6in
086        double tagWidth = Units.inchesToMeters(6);
087        TargetModel tagModel = TargetModel.kAprilTag16h5;
088
089        params.tagFamily =
090                switch (settings.tagFamily) {
091                    case kTag36h11 -> {
092                        // 2024 tag, 6.5in
093                        tagWidth = Units.inchesToMeters(6.5);
094                        tagModel = TargetModel.kAprilTag36h11;
095                        yield Objdetect.DICT_APRILTAG_36h11;
096                    }
097                    case kTag16h5 -> {
098                        // 2024 tag, 6.5in
099                        tagWidth = Units.inchesToMeters(6);
100                        tagModel = TargetModel.kAprilTag16h5;
101                        yield Objdetect.DICT_APRILTAG_16h5;
102                    }
103                };
104
105        int threshMinSize = Math.max(3, settings.threshWinSizes.getFirst());
106        settings.threshWinSizes.setFirst(threshMinSize);
107        params.threshMinSize = threshMinSize;
108        int threshStepSize = Math.max(2, settings.threshStepSize);
109        settings.threshStepSize = threshStepSize;
110        params.threshStepSize = threshStepSize;
111        int threshMaxSize = Math.max(threshMinSize, settings.threshWinSizes.getSecond());
112        settings.threshWinSizes.setSecond(threshMaxSize);
113        params.threshMaxSize = threshMaxSize;
114        params.threshConstant = settings.threshConstant;
115
116        params.useCornerRefinement = settings.useCornerRefinement;
117        params.refinementMaxIterations = settings.refineNumIterations;
118        params.refinementMinErrorPx = settings.refineMinErrorPx;
119        params.useAruco3 = settings.useAruco3;
120        params.aruco3MinMarkerSideRatio = settings.aruco3MinMarkerSideRatio;
121        params.aruco3MinCanonicalImgSide = settings.aruco3MinCanonicalImgSide;
122        arucoDetectionPipe.setParams(params);
123
124        if (frameStaticProperties.cameraCalibration != null) {
125            var cameraMatrix = frameStaticProperties.cameraCalibration.getCameraIntrinsicsMat();
126            if (cameraMatrix != null && cameraMatrix.rows() > 0) {
127                var estimatorParams =
128                        new ArucoPoseEstimatorPipeParams(frameStaticProperties.cameraCalibration, tagWidth);
129                singleTagPoseEstimatorPipe.setParams(estimatorParams);
130
131                // TODO global state ew
132                var atfl = ConfigManager.getInstance().getConfig().getApriltagFieldLayout();
133                multiTagPNPPipe.setParams(
134                        new MultiTargetPNPPipeParams(frameStaticProperties.cameraCalibration, atfl, tagModel));
135            }
136        }
137    }
138
139    @Override
140    protected CVPipelineResult process(Frame frame, ArucoPipelineSettings settings) {
141        long sumPipeNanosElapsed = 0L;
142
143        if (frame.type != FrameThresholdType.GREYSCALE) {
144            // We asked for a GREYSCALE frame, but didn't get one -- best we can do is give up
145            return new CVPipelineResult(frame.sequenceID, 0, 0, List.of(), frame);
146        }
147
148        CVPipeResult<List<ArucoDetectionResult>> tagDetectionPipeResult;
149        tagDetectionPipeResult = arucoDetectionPipe.run(frame.processedImage);
150        sumPipeNanosElapsed += tagDetectionPipeResult.nanosElapsed;
151
152        // If we want to debug the thresholding steps, draw the first step to the color image
153        if (settings.debugThreshold) {
154            drawThresholdFrame(
155                    frame.processedImage.getMat(),
156                    frame.colorImage.getMat(),
157                    settings.threshWinSizes.getFirst(),
158                    settings.threshConstant);
159        }
160
161        List<TrackedTarget> targetList = new ArrayList<>();
162        for (ArucoDetectionResult detection : tagDetectionPipeResult.output) {
163            // Populate target list for multitag
164            // (TODO: Address circular dependencies. Multitag only requires corners and IDs, this should
165            // not be necessary.)
166            TrackedTarget target =
167                    new TrackedTarget(
168                            detection,
169                            null,
170                            new TargetCalculationParameters(
171                                    false, null, null, null, null, frameStaticProperties));
172
173            targetList.add(target);
174        }
175
176        // Do multi-tag pose estimation
177        Optional<MultiTargetPNPResult> multiTagResult = Optional.empty();
178        if (settings.solvePNPEnabled && settings.doMultiTarget) {
179            var multiTagOutput = multiTagPNPPipe.run(targetList);
180            sumPipeNanosElapsed += multiTagOutput.nanosElapsed;
181            multiTagResult = multiTagOutput.output;
182        }
183
184        // Do single-tag pose estimation
185        if (settings.solvePNPEnabled) {
186            // Clear target list that was used for multitag so we can add target transforms
187            targetList.clear();
188            // TODO global state again ew
189            var atfl = ConfigManager.getInstance().getConfig().getApriltagFieldLayout();
190
191            for (ArucoDetectionResult detection : tagDetectionPipeResult.output) {
192                AprilTagPoseEstimate tagPoseEstimate = null;
193                // Do single-tag estimation when "always enabled" or if a tag was not used for multitag
194                if (settings.doSingleTargetAlways
195                        || !(multiTagResult.isPresent()
196                                && multiTagResult.get().fiducialIDsUsed.contains((short) detection.getId()))) {
197                    var poseResult = singleTagPoseEstimatorPipe.run(detection);
198                    sumPipeNanosElapsed += poseResult.nanosElapsed;
199                    tagPoseEstimate = poseResult.output;
200                }
201
202                // If single-tag estimation was not done, this is a multi-target tag from the layout
203                if (tagPoseEstimate == null && multiTagResult.isPresent()) {
204                    // compute this tag's camera-to-tag transform using the multitag result
205                    var tagPose = atfl.getTagPose(detection.getId());
206                    if (tagPose.isPresent()) {
207                        var camToTag =
208                                new Transform3d(
209                                        new Pose3d().plus(multiTagResult.get().estimatedPose.best), tagPose.get());
210                        // match expected OpenCV coordinate system
211                        camToTag =
212                                CoordinateSystem.convert(camToTag, CoordinateSystem.NWU(), CoordinateSystem.EDN());
213
214                        tagPoseEstimate = new AprilTagPoseEstimate(camToTag, camToTag, 0, 0);
215                    }
216                }
217
218                // populate the target list
219                // Challenge here is that TrackedTarget functions with OpenCV Contour
220                TrackedTarget target =
221                        new TrackedTarget(
222                                detection,
223                                tagPoseEstimate,
224                                new TargetCalculationParameters(
225                                        false, null, null, null, null, frameStaticProperties));
226
227                var correctedBestPose =
228                        MathUtils.convertOpenCVtoPhotonTransform(target.getBestCameraToTarget3d());
229                var correctedAltPose =
230                        MathUtils.convertOpenCVtoPhotonTransform(target.getAltCameraToTarget3d());
231
232                target.setBestCameraToTarget3d(
233                        new Transform3d(correctedBestPose.getTranslation(), correctedBestPose.getRotation()));
234                target.setAltCameraToTarget3d(
235                        new Transform3d(correctedAltPose.getTranslation(), correctedAltPose.getRotation()));
236
237                targetList.add(target);
238            }
239        }
240
241        var fpsResult = calculateFPSPipe.run(null);
242        var fps = fpsResult.output;
243
244        return new CVPipelineResult(
245                frame.sequenceID, sumPipeNanosElapsed, fps, targetList, multiTagResult, frame);
246    }
247
248    private void drawThresholdFrame(Mat greyMat, Mat outputMat, int windowSize, double constant) {
249        if (windowSize % 2 == 0) windowSize++;
250        Imgproc.adaptiveThreshold(
251                greyMat,
252                outputMat,
253                255,
254                Imgproc.ADAPTIVE_THRESH_MEAN_C,
255                Imgproc.THRESH_BINARY_INV,
256                windowSize,
257                constant);
258    }
259
260    @Override
261    public void release() {
262        arucoDetectionPipe.release();
263        singleTagPoseEstimatorPipe.release();
264        arucoDetectionPipe = null;
265        singleTagPoseEstimatorPipe = null;
266        super.release();
267    }
268}