USB.h

#include <cscore_cv.h>
 
#include <frc/apriltag/AprilTagDetector.h>
#include <frc/apriltag/AprilTagFieldLayout.h>
#include <frc/apriltag/AprilTagPoseEstimator.h>
 
#include "valkyrie/sensors/PoseEstimateSensor.h"
 
namespace valor {
namespace sensors {
namespace apriltag {
 
class USB : public PoseEstimateSensor {
   public:
    enum class Solver {
        SOLVE_PNP,
        BEST_SCORE
    };

    struct Config {
        units::meter_t fieldBorderMargin = 0.5_m;

        units::meter_t translationalStdDevCoefficient = 0.01_m;
        units::radian_t rotationalStdDevCoefficient = 0.03_rad;
        double tagWeight = 1;

        double ambiguityThreshold = 0.4;
 
        Solver solver = Solver::BEST_SCORE;
    };
 
    using VisionFilter = std::function<bool(std::span<frc::AprilTagDetection const* const> detections, frc::Pose3d estimate,
                                            const Config& config, const frc::AprilTagFieldLayout& fieldLayout)>;
 
    using ScoreCalculator = std::function<double(const frc::AprilTagDetection& detection, frc::Pose3d estimate, const Config& config,
                                                 const frc::AprilTagFieldLayout& fieldLayout)>;
 
    using DoubtCalculator =
        std::function<PoseEstimate::Doubts(std::span<frc::AprilTagDetection const* const> detections, frc::Pose3d estimate, Config& config,
                                           const frc::AprilTagFieldLayout& fieldLayout)>;
 
    USB(std::shared_ptr<cs::UsbCamera> camera, frc::AprilTagPoseEstimator::Config config,
        frc::AprilTagField field = frc::AprilTagField::kDefaultField);
 
    constexpr const Config& GetConfig() const { return config; }
 
    constexpr void SetConfig(const Config& c) { config = c; }
 
    inline void SetVisionFilter(VisionFilter f) { visionFilter = f; }
 
    inline void SetScoreCalculator(ScoreCalculator s) { scoreCalculator = s; }
 
    inline void SetDoubtCalculator(DoubtCalculator d) { doubtCalculator = d; }
 
    constexpr std::vector<PoseEstimate> GetAllEstimates() const { return estimates; }
 
    static frc::AprilTagDetector& GetDetector();
 
    static bool StandardVisionFilter(std::span<frc::AprilTagDetection const* const> detections, frc::Pose3d estimate, const Config& config,
                                     const frc::AprilTagFieldLayout& fieldLayout);
 
    static double StandardScoreCalculator(const frc::AprilTagDetection& detection, frc::Pose3d estimate, const Config& config,
                                          const frc::AprilTagFieldLayout& fieldLayout);
 
    static PoseEstimate::Doubts StandardDoubtCalculator(std::span<frc::AprilTagDetection const* const> detections, frc::Pose3d estimate,
                                                        const Config& config, const frc::AprilTagFieldLayout& fieldLayout);
 
   private:
    bool ConvertGrayMat();
    void DrawTag(const frc::AprilTagDetection& detection);
    std::optional<PoseEstimate> CalculateHighestScore(std::span<frc::AprilTagDetection const* const>);
    std::optional<PoseEstimate> CalculatePnP(std::span<frc::AprilTagDetection const* const>);
    std::optional<PoseEstimate> GetEstimate();
 
    Config config;
    std::vector<PoseEstimate> estimates;
 
    cs::CvSink sink;
    cs::CvSource outputStream;
 
    cv::Mat mat;
    cv::Mat grayMat;
    cv::Mat cameraMatrix;
    cv::Mat distortionCoefficients = cv::Mat::zeros(4, 1, CV_64F);
    cv::Mat rvec;
    cv::Mat tvec;
    cv::Mat R;
    std::vector<cv::Point3d> objectPoints;
    std::vector<cv::Point2d> imagePoints;
 
    std::shared_ptr<cs::UsbCamera> camera;
    frc::AprilTagPoseEstimator estimator;
    frc::AprilTagFieldLayout layout;
 
    VisionFilter visionFilter = StandardVisionFilter;
    ScoreCalculator scoreCalculator = StandardScoreCalculator;
    DoubtCalculator doubtCalculator = StandardDoubtCalculator;
};
 
}  // namespace apriltag
}  // namespace sensors
}  // namespace valor