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.estimation;
019
020import edu.wpi.first.apriltag.AprilTag;
021import edu.wpi.first.apriltag.AprilTagFieldLayout;
022import edu.wpi.first.cscore.OpenCvLoader;
023import edu.wpi.first.math.MatBuilder;
024import edu.wpi.first.math.Matrix;
025import edu.wpi.first.math.Nat;
026import edu.wpi.first.math.geometry.Pose3d;
027import edu.wpi.first.math.geometry.Rotation2d;
028import edu.wpi.first.math.geometry.Transform2d;
029import edu.wpi.first.math.geometry.Transform3d;
030import edu.wpi.first.math.geometry.Translation3d;
031import edu.wpi.first.math.numbers.*;
032import java.util.ArrayList;
033import java.util.List;
034import java.util.Objects;
035import java.util.Optional;
036import java.util.stream.Collectors;
037import org.ejml.simple.SimpleMatrix;
038import org.opencv.calib3d.Calib3d;
039import org.opencv.core.MatOfDouble;
040import org.opencv.core.MatOfPoint2f;
041import org.opencv.core.Point;
042import org.photonvision.jni.ConstrainedSolvepnpJni;
043import org.photonvision.targeting.PhotonTrackedTarget;
044import org.photonvision.targeting.PnpResult;
045import org.photonvision.targeting.TargetCorner;
046
047public class VisionEstimation {
048    private VisionEstimation() {}
049
050    /**
051     * Get the list of visible {@link AprilTag}s which are in the tag layout using the visible tag
052     * IDs.
053     *
054     * @param visTags The list of targets to search for visible tags.
055     * @param tagLayout The tag layout to search
056     */
057    public static List<AprilTag> getVisibleLayoutTags(
058            List<PhotonTrackedTarget> visTags, AprilTagFieldLayout tagLayout) {
059        return visTags.stream()
060                .map(
061                        t -> {
062                            int id = t.getFiducialId();
063                            var maybePose = tagLayout.getTagPose(id);
064                            return maybePose.map(pose3d -> new AprilTag(id, pose3d)).orElse(null);
065                        })
066                .filter(Objects::nonNull)
067                .collect(Collectors.toList());
068    }
069
070    /**
071     * Performs solvePNP using 3d-2d point correspondences of visible AprilTags to estimate the
072     * field-to-camera transformation. If only one tag is visible, the result may have an alternate
073     * solution.
074     *
075     * <p><b>Note:</b> The returned transformation is from the field origin to the camera pose!
076     *
077     * <p>With only one tag: {@link OpenCVHelp#solvePNP_SQUARE}
078     *
079     * <p>With multiple tags: {@link OpenCVHelp#solvePNP_SQPNP}
080     *
081     * @param cameraMatrix The camera intrinsics matrix in standard opencv form
082     * @param distCoeffs The camera distortion matrix in standard opencv form
083     * @param visTags The visible tags reported by PV. Non-tag targets are automatically excluded.
084     * @param tagLayout The known tag layout on the field
085     * @param tagModel The model describing the tag's geometry
086     * @return The transformation that maps the field origin to the camera pose. Ensure the {@link
087     *     PnpResult} are present before utilizing them.
088     */
089    public static Optional<PnpResult> estimateCamPosePNP(
090            Matrix<N3, N3> cameraMatrix,
091            Matrix<N8, N1> distCoeffs,
092            List<PhotonTrackedTarget> visTags,
093            AprilTagFieldLayout tagLayout,
094            TargetModel tagModel) {
095        if (tagLayout == null
096                || visTags == null
097                || tagLayout.getTags().isEmpty()
098                || visTags.isEmpty()) {
099            return Optional.empty();
100        }
101
102        var corners = new ArrayList<TargetCorner>();
103        var knownTags = new ArrayList<AprilTag>();
104        // ensure these are AprilTags in our layout
105        for (var tgt : visTags) {
106            int id = tgt.getFiducialId();
107            tagLayout
108                    .getTagPose(id)
109                    .ifPresent(
110                            pose -> {
111                                knownTags.add(new AprilTag(id, pose));
112                                corners.addAll(tgt.getDetectedCorners());
113                            });
114        }
115        if (knownTags.isEmpty() || corners.isEmpty() || corners.size() % 4 != 0) {
116            return Optional.empty();
117        }
118        OpenCvLoader.forceStaticLoad();
119
120        Point[] points = OpenCVHelp.cornersToPoints(corners);
121
122        // single-tag pnp
123        if (knownTags.size() == 1) {
124            var camToTag =
125                    OpenCVHelp.solvePNP_SQUARE(cameraMatrix, distCoeffs, tagModel.vertices, points);
126            if (!camToTag.isPresent()) return Optional.empty();
127            var bestPose = knownTags.get(0).pose.transformBy(camToTag.get().best.inverse());
128            var altPose = new Pose3d();
129            if (camToTag.get().ambiguity != 0)
130                altPose = knownTags.get(0).pose.transformBy(camToTag.get().alt.inverse());
131
132            var o = new Pose3d();
133            return Optional.of(
134                    new PnpResult(
135                            new Transform3d(o, bestPose),
136                            new Transform3d(o, altPose),
137                            camToTag.get().ambiguity,
138                            camToTag.get().bestReprojErr,
139                            camToTag.get().altReprojErr));
140        }
141        // multi-tag pnp
142        else {
143            var objectTrls = new ArrayList<Translation3d>();
144            for (var tag : knownTags) objectTrls.addAll(tagModel.getFieldVertices(tag.pose));
145            var camToOrigin = OpenCVHelp.solvePNP_SQPNP(cameraMatrix, distCoeffs, objectTrls, points);
146            if (camToOrigin.isEmpty()) return Optional.empty();
147            return Optional.of(
148                    new PnpResult(
149                            camToOrigin.get().best.inverse(),
150                            camToOrigin.get().alt.inverse(),
151                            camToOrigin.get().ambiguity,
152                            camToOrigin.get().bestReprojErr,
153                            camToOrigin.get().altReprojErr));
154        }
155    }
156
157    /**
158     * Performs constrained solvePNP using 3d-2d point correspondences of visible AprilTags to
159     * estimate the field-to-camera transformation.
160     *
161     * <p><b>Note:</b> The returned transformation is from the field origin to the robot drivebase
162     *
163     * @param cameraMatrix The camera intrinsics matrix in standard opencv form
164     * @param distCoeffs The camera distortion matrix in standard opencv form
165     * @param visTags The visible tags reported by PV. Non-tag targets are automatically excluded.
166     * @param robot2camera The {@link Transform3d} from the robot odometry frame to the camera optical
167     *     frame
168     * @param robotPoseSeed An initial guess at robot pose, refined via optimization. Better guesses
169     *     will converge faster.
170     * @param tagLayout The known tag layout on the field
171     * @param tagModel The physical size of the AprilTags
172     * @param headingFree If heading is completely free, or if our measured gyroθ is taken into
173     *     account
174     * @param gyroθ If headingFree is false, the best estimate at robot yaw. Excursions from this are
175     *     penalized in our cost function.
176     * @param gyroErrorScaleFac A relative weight for gyro heading excursions against tag corner
177     *     reprojection error.
178     * @return The transformation that maps the field origin to the camera pose. Ensure the {@link
179     *     PnpResult} are present before utilizing them.
180     */
181    public static Optional<PnpResult> estimateRobotPoseConstrainedSolvepnp(
182            Matrix<N3, N3> cameraMatrix,
183            Matrix<N8, N1> distCoeffs,
184            List<PhotonTrackedTarget> visTags,
185            Transform3d robot2camera,
186            Pose3d robotPoseSeed,
187            AprilTagFieldLayout tagLayout,
188            TargetModel tagModel,
189            boolean headingFree,
190            Rotation2d gyroθ,
191            double gyroErrorScaleFac) {
192        if (tagLayout == null
193                || visTags == null
194                || tagLayout.getTags().isEmpty()
195                || visTags.isEmpty()) {
196            return Optional.empty();
197        }
198
199        var corners = new ArrayList<TargetCorner>();
200        var knownTags = new ArrayList<AprilTag>();
201        // ensure these are AprilTags in our layout
202        for (var tgt : visTags) {
203            int id = tgt.getFiducialId();
204            tagLayout
205                    .getTagPose(id)
206                    .ifPresent(
207                            pose -> {
208                                knownTags.add(new AprilTag(id, pose));
209                                corners.addAll(tgt.getDetectedCorners());
210                            });
211        }
212        if (knownTags.isEmpty() || corners.isEmpty() || corners.size() % 4 != 0) {
213            return Optional.empty();
214        }
215        OpenCvLoader.forceStaticLoad();
216
217        Point[] points = OpenCVHelp.cornersToPoints(corners);
218
219        // Undistort
220        {
221            MatOfPoint2f temp = new MatOfPoint2f();
222            MatOfDouble cameraMatrixMat = new MatOfDouble();
223            MatOfDouble distCoeffsMat = new MatOfDouble();
224            OpenCVHelp.matrixToMat(cameraMatrix.getStorage()).assignTo(cameraMatrixMat);
225            OpenCVHelp.matrixToMat(distCoeffs.getStorage()).assignTo(distCoeffsMat);
226
227            temp.fromArray(points);
228            Calib3d.undistortImagePoints(temp, temp, cameraMatrixMat, distCoeffsMat);
229            points = temp.toArray();
230
231            temp.release();
232            cameraMatrixMat.release();
233            distCoeffsMat.release();
234        }
235
236        // Rotate from wpilib to opencv camera CS
237        var robot2cameraBase =
238                MatBuilder.fill(Nat.N4(), Nat.N4(), 0, 0, 1, 0, -1, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 1);
239        var robotToCamera = robot2camera.toMatrix().times(robot2cameraBase);
240
241        // Where we saw corners
242        var point_observations = new SimpleMatrix(2, points.length);
243        for (int i = 0; i < points.length; i++) {
244            point_observations.set(0, i, points[i].x);
245            point_observations.set(1, i, points[i].y);
246        }
247
248        // Affine corner locations
249        var objectTrls = new ArrayList<Translation3d>();
250        for (var tag : knownTags) objectTrls.addAll(tagModel.getFieldVertices(tag.pose));
251        var field2points = new SimpleMatrix(4, points.length);
252        for (int i = 0; i < objectTrls.size(); i++) {
253            field2points.set(0, i, objectTrls.get(i).getX());
254            field2points.set(1, i, objectTrls.get(i).getY());
255            field2points.set(2, i, objectTrls.get(i).getZ());
256            field2points.set(3, i, 1);
257        }
258
259        // fx fy cx cy
260        double[] cameraCal = {
261            cameraMatrix.get(0, 0),
262            cameraMatrix.get(1, 1),
263            cameraMatrix.get(0, 2),
264            cameraMatrix.get(1, 2),
265        };
266
267        var guess2 = robotPoseSeed.toPose2d();
268
269        var ret =
270                ConstrainedSolvepnpJni.do_optimization(
271                        headingFree,
272                        knownTags.size(),
273                        cameraCal,
274                        robotToCamera.getData(),
275                        new double[] {
276                            guess2.getX(), guess2.getY(), guess2.getRotation().getRadians(),
277                        },
278                        field2points.getDDRM().getData(),
279                        point_observations.getDDRM().getData(),
280                        gyroθ.getRadians(),
281                        gyroErrorScaleFac);
282
283        if (ret == null) {
284            return Optional.empty();
285        } else {
286            var pnpresult = new PnpResult();
287            pnpresult.best = new Transform3d(new Transform2d(ret[0], ret[1], new Rotation2d(ret[2])));
288            return Optional.of(pnpresult);
289        }
290    }
291}