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
018package org.photonvision.vision.pipe.impl;
019
020import edu.wpi.first.apriltag.AprilTagPoseEstimate;
021import edu.wpi.first.math.VecBuilder;
022import edu.wpi.first.math.geometry.Rotation3d;
023import edu.wpi.first.math.geometry.Transform3d;
024import edu.wpi.first.math.geometry.Translation3d;
025import java.util.ArrayList;
026import java.util.List;
027import org.opencv.calib3d.Calib3d;
028import org.opencv.core.CvType;
029import org.opencv.core.Mat;
030import org.opencv.core.MatOfPoint2f;
031import org.opencv.core.MatOfPoint3f;
032import org.opencv.core.Point3;
033import org.photonvision.vision.aruco.ArucoDetectionResult;
034import org.photonvision.vision.calibration.CameraCalibrationCoefficients;
035import org.photonvision.vision.opencv.Releasable;
036import org.photonvision.vision.pipe.CVPipe;
037
038public class ArucoPoseEstimatorPipe
039        extends CVPipe<
040                ArucoDetectionResult,
041                AprilTagPoseEstimate,
042                ArucoPoseEstimatorPipe.ArucoPoseEstimatorPipeParams>
043        implements Releasable {
044    // image points of marker corners
045    private final MatOfPoint2f imagePoints = new MatOfPoint2f(Mat.zeros(4, 1, CvType.CV_32FC2));
046    // rvec/tvec estimations from solvepnp
047    private final List<Mat> rvecs = new ArrayList<>();
048    private final List<Mat> tvecs = new ArrayList<>();
049    // unused parameters
050    private final Mat rvec = Mat.zeros(3, 1, CvType.CV_32F);
051    private final Mat tvec = Mat.zeros(3, 1, CvType.CV_32F);
052    // reprojection error of solvepnp estimations
053    private final Mat reprojectionErrors = Mat.zeros(2, 1, CvType.CV_32F);
054
055    // Tag corner locations in object space - order matters for ippe_square
056    MatOfPoint3f objectPoints = new MatOfPoint3f();
057
058    private final int kNaNRetries = 1;
059
060    private Translation3d tvecToTranslation3d(Mat mat) {
061        double[] tvec = new double[3];
062        mat.get(0, 0, tvec);
063        return new Translation3d(tvec[0], tvec[1], tvec[2]);
064    }
065
066    private Rotation3d rvecToRotation3d(Mat mat) {
067        double[] rvec = new double[3];
068        mat.get(0, 0, rvec);
069        return new Rotation3d(VecBuilder.fill(rvec[0], rvec[1], rvec[2]));
070    }
071
072    @Override
073    protected AprilTagPoseEstimate process(ArucoDetectionResult in) {
074        // We receive 2d corners as (BL, BR, TR, TL) but we want (BR, BL, TL, TR)
075        double[] xCorn = in.getXCorners();
076        double[] yCorn = in.getYCorners();
077        imagePoints.put(0, 0, new float[] {(float) xCorn[1], (float) yCorn[1]});
078        imagePoints.put(1, 0, new float[] {(float) xCorn[0], (float) yCorn[0]});
079        imagePoints.put(2, 0, new float[] {(float) xCorn[3], (float) yCorn[3]});
080        imagePoints.put(3, 0, new float[] {(float) xCorn[2], (float) yCorn[2]});
081
082        float[] reprojErrors = new float[2];
083        // very rarely the solvepnp solver returns NaN results, so we retry with slight noise added
084        for (int i = 0; i < kNaNRetries + 1; i++) {
085            // SolvePnP with SOLVEPNP_IPPE_SQUARE solver
086            Calib3d.solvePnPGeneric(
087                    objectPoints,
088                    imagePoints,
089                    params.calibration.getCameraIntrinsicsMat(),
090                    params.calibration.getDistCoeffsMat(),
091                    rvecs,
092                    tvecs,
093                    false,
094                    Calib3d.SOLVEPNP_IPPE_SQUARE,
095                    rvec,
096                    tvec,
097                    reprojectionErrors);
098
099            // check if we got a NaN result
100            reprojectionErrors.get(0, 0, reprojErrors);
101            if (!Double.isNaN(reprojErrors[0])) break;
102            else { // add noise and retry
103                double[] br = imagePoints.get(0, 0);
104                br[0] -= 0.001;
105                br[1] -= 0.001;
106                imagePoints.put(0, 0, br);
107            }
108        }
109
110        // create AprilTagPoseEstimate with results
111        if (tvecs.isEmpty())
112            return new AprilTagPoseEstimate(new Transform3d(), new Transform3d(), 0, 0);
113        return new AprilTagPoseEstimate(
114                new Transform3d(tvecToTranslation3d(tvecs.get(0)), rvecToRotation3d(rvecs.get(0))),
115                new Transform3d(tvecToTranslation3d(tvecs.get(1)), rvecToRotation3d(rvecs.get(1))),
116                reprojErrors[0],
117                reprojErrors[1]);
118    }
119
120    @Override
121    public void setParams(ArucoPoseEstimatorPipe.ArucoPoseEstimatorPipeParams newParams) {
122        // exact equality check OK here, the number shouldn't change
123        if (this.params == null || this.params.tagSize != newParams.tagSize) {
124            var tagSize = newParams.tagSize;
125
126            // This order is relevant for SOLVEPNP_IPPE_SQUARE
127            // The expected 2d correspondences with a tag facing the camera would be (BR, BL, TL, TR)
128            objectPoints.fromArray(
129                    new Point3(-tagSize / 2, tagSize / 2, 0),
130                    new Point3(tagSize / 2, tagSize / 2, 0),
131                    new Point3(tagSize / 2, -tagSize / 2, 0),
132                    new Point3(-tagSize / 2, -tagSize / 2, 0));
133        }
134
135        super.setParams(newParams);
136    }
137
138    @Override
139    public void release() {
140        imagePoints.release();
141        for (var m : rvecs) m.release();
142        rvecs.clear();
143        for (var m : tvecs) m.release();
144        tvecs.clear();
145        rvec.release();
146        tvec.release();
147        reprojectionErrors.release();
148    }
149
150    public static class ArucoPoseEstimatorPipeParams {
151        final CameraCalibrationCoefficients calibration;
152        final double tagSize;
153
154        // object vertices defined by tag size
155
156        public ArucoPoseEstimatorPipeParams(CameraCalibrationCoefficients cal, double tagSize) {
157            this.calibration = cal;
158            this.tagSize = tagSize;
159        }
160    }
161}