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