Extract common charuco detection code into charuco_lib
This lets us reuse all the charuco extraction code when we do extrinsics
calibration soon. We really just want a callback with the detected
board.
Change-Id: If806b026cc67bf9e7b61935f93b1902acdcb2783
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/y2020/vision/charuco_lib.h b/y2020/vision/charuco_lib.h
new file mode 100644
index 0000000..e8dc3eb
--- /dev/null
+++ b/y2020/vision/charuco_lib.h
@@ -0,0 +1,112 @@
+#ifndef Y2020_VISION_CHARUCO_LIB_H_
+#define Y2020_VISION_CHARUCO_LIB_H_
+
+#include <functional>
+#include <string_view>
+
+#include <opencv2/aruco/charuco.hpp>
+#include <opencv2/calib3d.hpp>
+#include "Eigen/Dense"
+#include "Eigen/Geometry"
+
+#include "absl/types/span.h"
+#include "aos/events/event_loop.h"
+#include "y2020/vision/sift/sift_generated.h"
+#include "y2020/vision/sift/sift_training_generated.h"
+
+namespace frc971 {
+namespace vision {
+
+// Class to find extrinsics for a specified pi's camera using the provided
+// training data.
+class CameraCalibration {
+ public:
+ CameraCalibration(const absl::Span<const uint8_t> training_data_bfbs,
+ std::string_view pi);
+
+ // Intrinsics for the located camera.
+ cv::Mat CameraIntrinsics() const;
+ Eigen::Matrix3d CameraIntrinsicsEigen() const;
+
+ // Distortion coefficients for the located camera.
+ cv::Mat CameraDistCoeffs() const;
+
+ private:
+ // Finds the camera specific calibration flatbuffer.
+ const sift::CameraCalibration *FindCameraCalibration(
+ const sift::TrainingData *const training_data, std::string_view pi) const;
+
+ // Pointer to this camera's calibration parameters.
+ const sift::CameraCalibration *camera_calibration_;
+};
+
+// Class to call a function with a cv::Mat and age when an image shows up on the
+// provided channel. This hides all the conversions and wrangling needed to
+// view the image.
+class ImageCallback {
+ public:
+ ImageCallback(aos::EventLoop *event_loop, std::string_view channel,
+ std::function<void(cv::Mat, double)> &&fn);
+
+ private:
+ aos::EventLoop *event_loop_;
+ std::function<void(cv::Mat, double)> handle_image_;
+};
+
+// Class which calls a callback each time an image arrives with the information
+// extracted from it.
+class CharucoExtractor {
+ public:
+ // The callback takes the following arguments:
+ // cv::Mat -> image with overlays drawn on it.
+ // const double -> Duration between when the image was captured and this
+ // callback was called.
+ // std::vector<int> -> charuco_ids
+ // std::vector<cv::Point2f> -> charuco_corners
+ // bool -> true if rvec/tvec is valid.
+ // Eigen::Vector3d -> rvec
+ // Eigen::Vector3d -> tvec
+ CharucoExtractor(aos::EventLoop *event_loop, std::string_view pi,
+ std::function<void(cv::Mat, const double, std::vector<int>,
+ std::vector<cv::Point2f>, bool,
+ Eigen::Vector3d, Eigen::Vector3d)> &&fn);
+
+ // Returns the aruco dictionary in use.
+ cv::Ptr<cv::aruco::Dictionary> dictionary() const { return dictionary_; }
+ // Returns the aruco board in use.
+ cv::Ptr<cv::aruco::CharucoBoard> board() const { return board_; }
+
+ // Returns the camera matrix for this camera.
+ const cv::Mat camera_matrix() const { return camera_matrix_; }
+ // Returns the distortion coefficients for this camera.
+ const cv::Mat dist_coeffs() const { return dist_coeffs_; }
+
+ private:
+ // Handles the image by detecting the charuco board in it and calling the
+ // callback with the extracted board corners, corresponding IDs, and pose.
+ void HandleImage(cv::Mat rgb_image, const double age_double);
+
+ CameraCalibration calibration_;
+
+ cv::Ptr<cv::aruco::Dictionary> dictionary_;
+ cv::Ptr<cv::aruco::CharucoBoard> board_;
+
+ const cv::Mat camera_matrix_;
+ const Eigen::Matrix3d eigen_camera_matrix_;
+ const cv::Mat dist_coeffs_;
+
+ const std::optional<uint16_t> pi_number_;
+
+ ImageCallback image_callback_;
+
+ // Function to call.
+ std::function<void(cv::Mat, const double, std::vector<int>,
+ std::vector<cv::Point2f>, bool, Eigen::Vector3d,
+ Eigen::Vector3d)>
+ handle_charuco_;
+};
+
+} // namespace vision
+} // namespace frc971
+
+#endif // Y2020_VISION_CHARUCO_LIB_H_