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.math.VecBuilder;
021import edu.wpi.first.math.geometry.Rotation3d;
022import edu.wpi.first.math.geometry.Transform3d;
023import edu.wpi.first.math.geometry.Translation3d;
024import java.util.List;
025import org.opencv.calib3d.Calib3d;
026import org.opencv.core.Core;
027import org.opencv.core.Mat;
028import org.opencv.core.MatOfPoint2f;
029import org.photonvision.common.logging.LogGroup;
030import org.photonvision.common.logging.Logger;
031import org.photonvision.common.util.math.MathUtils;
032import org.photonvision.vision.calibration.CameraCalibrationCoefficients;
033import org.photonvision.vision.pipe.CVPipe;
034import org.photonvision.vision.target.TargetModel;
035import org.photonvision.vision.target.TrackedTarget;
036
037public class SolvePNPPipe
038        extends CVPipe<List<TrackedTarget>, List<TrackedTarget>, SolvePNPPipe.SolvePNPPipeParams> {
039    private static final Logger logger = new Logger(SolvePNPPipe.class, LogGroup.VisionModule);
040
041    private final MatOfPoint2f imagePoints = new MatOfPoint2f();
042
043    private boolean hasWarned = false;
044
045    @Override
046    protected List<TrackedTarget> process(List<TrackedTarget> targetList) {
047        if (params.cameraCoefficients == null
048                || params.cameraCoefficients.getCameraIntrinsicsMat() == null
049                || params.cameraCoefficients.getDistCoeffsMat() == null) {
050            if (!hasWarned) {
051                logger.warn(
052                        "Cannot perform solvePNP an uncalibrated camera! Please calibrate this resolution...");
053                hasWarned = true;
054            }
055            return targetList;
056        }
057
058        for (var target : targetList) {
059            calculateTargetPose(target);
060        }
061        return targetList;
062    }
063
064    private void calculateTargetPose(TrackedTarget target) {
065        var corners = target.getTargetCorners();
066        if (corners == null
067                || corners.isEmpty()
068                || params.cameraCoefficients == null
069                || params.cameraCoefficients.getCameraIntrinsicsMat() == null
070                || params.cameraCoefficients.getDistCoeffsMat() == null) {
071            return;
072        }
073        this.imagePoints.fromList(corners);
074
075        var rVec = new Mat();
076        var tVec = new Mat();
077        try {
078            Calib3d.solvePnP(
079                    params.targetModel.getRealWorldTargetCoordinates(),
080                    imagePoints,
081                    params.cameraCoefficients.getCameraIntrinsicsMat(),
082                    params.cameraCoefficients.getDistCoeffsMat(),
083                    rVec,
084                    tVec);
085        } catch (Exception e) {
086            logger.error("Exception when attempting solvePnP!", e);
087            return;
088        }
089
090        target.setCameraRelativeTvec(tVec);
091        target.setCameraRelativeRvec(rVec);
092
093        Translation3d translation =
094                new Translation3d(tVec.get(0, 0)[0], tVec.get(1, 0)[0], tVec.get(2, 0)[0]);
095        Rotation3d rotation =
096                new Rotation3d(
097                        VecBuilder.fill(rVec.get(0, 0)[0], rVec.get(1, 0)[0], rVec.get(2, 0)[0]),
098                        Core.norm(rVec));
099
100        Transform3d camToTarget =
101                MathUtils.convertOpenCVtoPhotonTransform(new Transform3d(translation, rotation));
102        target.setBestCameraToTarget3d(camToTarget);
103        target.setAltCameraToTarget3d(new Transform3d());
104    }
105
106    public static class SolvePNPPipeParams {
107        private final CameraCalibrationCoefficients cameraCoefficients;
108        private final TargetModel targetModel;
109
110        public SolvePNPPipeParams(
111                CameraCalibrationCoefficients cameraCoefficients, TargetModel targetModel) {
112            this.cameraCoefficients = cameraCoefficients;
113            this.targetModel = targetModel;
114        }
115    }
116}