Add nvidia apriltag backend

CMakeLists.txt

@@ -28,6 +28,7 @@ find_package(wpilib REQUIRED)
 find_package(OpenCV REQUIRED)
 find_package(Eigen3 REQUIRED)
 find_package(realsense2 REQUIRED)
+find_package(vpi REQUIRED)
 
 add_subdirectory(third_party/json)
 add_subdirectory(third_party/971apriltag)

src/fiddler_main.cc

@@ -7,37 +7,19 @@
 #include <thread>
 #include "apriltag/apriltag.h"
 #include "localization/position_sender.h"
-#include "localization/tag_estimator.h"
 #include "src/camera/camera_constants.h"
 #include "src/camera/camera_source.h"
 #include "src/camera/cscore_streamer.h"
 #include "src/camera/cv_camera.h"
+#include "src/localization/get_field_relitive_position.h"
+#include "src/localization/gpu_apriltag_detector.h"
+#include "src/localization/run_localization.h"
+#include "src/utils/camera_utils.h"
 #include "src/utils/nt_utils.h"
 #include "src/utils/timer.h"
 
 using json = nlohmann::json;
 
-void run_estimator(const int frame_width, const int frame_height,
-                   camera::CameraSource& source, std::string intrinsics,
-                   std::string extrinsics, uint port) {
-
-  localization::TagEstimator tag_estimator(frame_width, frame_height,
-                                           intrinsics, extrinsics);
-  localization::PositionSender position_sender(source.GetName());
-
-  camera::CscoreStreamer streamer(source.GetName(), port, 30, 1080, 1080);
-
-  while (true) {
-    utils::Timer timer(source.GetName(), false);
-    camera::timestamped_frame_t timestamped_frame = source.Get();
-    streamer.WriteFrame(timestamped_frame.frame);
-    std::vector<localization::tag_detection_t> estimates =
-        tag_estimator.Estimate(timestamped_frame.frame,
-                               timestamped_frame.timestamp);
-    position_sender.Send(estimates, timer.Stop());
-  }
-}
-
 int main() {
   utils::StartNetworktables();
 
@@ -52,16 +34,24 @@ int main() {
           camera::camera_constants[camera::Camera::USB1].pipeline)));
 
   std::thread usb0_thread(
-      run_estimator, 640, 480, std::ref(back_left_camera),
-      camera::camera_constants[camera::Camera::USB0].intrinsics_path,
-      camera::camera_constants[camera::Camera::USB0].extrinsics_path, 4971);
+      localization::run_localization, std::ref(back_left_camera),
+      std::make_unique<localization::GPUAprilTagDetector>(
+          640, 480,
+          utils::read_intrinsics(
+              camera::camera_constants[camera::Camera::USB0].intrinsics_path)),
+      camera::camera_constants[camera::Camera::USB0].extrinsics_path, 4971,
+      false);
 
   std::thread usb1_thread(
-      run_estimator, 1280, 720, std::ref(back_right_camera),
-      camera::camera_constants[camera::Camera::USB1].intrinsics_path,
-      camera::camera_constants[camera::Camera::USB1].extrinsics_path, 4972);
-
-  usb1_thread.join();
+      localization::run_localization, std::ref(back_right_camera),
+      std::make_unique<localization::GPUAprilTagDetector>(
+          1280, 720,
+          utils::read_intrinsics(
+              camera::camera_constants[camera::Camera::USB1].intrinsics_path)),
+      camera::camera_constants[camera::Camera::USB1].extrinsics_path, 4972,
+      false);
+
+  usb0_thread.join();
 
   return 0;
 }

src/localization/CMakeLists.txt

@@ -1,4 +1,4 @@
 cmake_minimum_required(VERSION 3.10)
 
-add_library(localization tag_estimator.cc position_sender.cc)
-target_link_libraries(localization PUBLIC ${OpenCV_LIBS} 971apriltag ntcore apriltag nlohmann_json::nlohmann_json wpilibc utils)
+add_library(localization position_sender.cc  gpu_apriltag_detector.cc get_field_relitive_position.cc run_localization.cc nvidia_apriltag_detector.cc)
+target_link_libraries(localization PUBLIC ${OpenCV_LIBS} 971apriltag ntcore apriltag nlohmann_json::nlohmann_json wpilibc utils camera)

src/localization/apriltag_detector.h

@@ -0,0 +1,22 @@
+#pragma once
+#include <frc/geometry/Transform3d.h>
+#include <opencv2/core/mat.hpp>
+#include "src/camera/camera_source.h"
+#include "src/localization/position.h"
+
+namespace localization {
+
+constexpr double ktag_size = 0.1651;  // meters
+const std::vector<cv::Point3f> kapriltag_dimensions = {
+    {-ktag_size / 2, ktag_size / 2, 0},
+    {ktag_size / 2, ktag_size / 2, 0},
+    {ktag_size / 2, -ktag_size / 2, 0},
+    {-ktag_size / 2, -ktag_size / 2, 0}};
+
+class IAprilTagDetector {
+ public:
+  virtual std::vector<tag_detection_t> GetTagDetections(
+      camera::timestamped_frame_t& frame) = 0;
+  virtual ~IAprilTagDetector() = default;
+};
+}  // namespace localization

src/localization/get_field_relitive_position.cc

@@ -0,0 +1,84 @@
+#include "src/localization/get_field_relitive_position.h"
+#include <frc/apriltag/AprilTagFieldLayout.h>
+#include <iomanip>
+#include <iostream>
+#include <nlohmann/json.hpp>
+#include "src/utils/log.h"
+namespace localization {
+
+tag_detection_t ToFeildRelitivePosition(
+    tag_detection_t tag_relative_position, frc::Transform3d camera_to_robot,
+    frc::AprilTagFieldLayout apriltag_layout, bool verbose) {
+
+  frc::Transform3d camera_to_tag(
+      units::meter_t{tag_relative_position.pose.X()},
+      units::meter_t{-tag_relative_position.pose.Y()},
+      units::meter_t{-tag_relative_position.pose.Z()},
+      frc::Rotation3d(
+          units::radian_t{tag_relative_position.pose.Rotation().X()},
+          units::radian_t{-tag_relative_position.pose.Rotation().Y()},
+          units::radian_t{-tag_relative_position.pose.Rotation().Z()} +
+              180_deg));
+
+  frc::Transform3d tag_to_camera = camera_to_tag.Inverse();
+
+  if (verbose) {
+    std::cout << "tag to camera: \n";
+    PrintTransform3d(tag_to_camera);
+    std::cout << "\n\n";
+  }
+
+  frc::Pose3d tag_pose =
+      apriltag_layout.GetTagPose(tag_relative_position.tag_id).value();
+
+  if (verbose) {
+    std::cout << "tag id: " << tag_relative_position.tag_id << std::endl;
+    std::cout << "tagpose: \n";
+    PrintPose3d(tag_pose);
+    std::cout << "\n\n";
+  }
+
+  frc::Pose3d camera_pose = tag_pose.TransformBy(tag_to_camera);
+
+  if (verbose) {
+    std::cout << "camerapose: \n";
+    PrintPose3d(camera_pose);
+    std::cout << "\n\n";
+  }
+
+  frc::Pose3d robot_pose = camera_pose.TransformBy(camera_to_robot);
+
+  tag_detection_t field_relative_pose;
+
+  field_relative_pose.tag_id = tag_relative_position.tag_id;
+
+  field_relative_pose.pose = robot_pose;
+
+  field_relative_pose.distance = tag_relative_position.distance;
+
+  field_relative_pose.timestamp = tag_relative_position.timestamp;
+
+  return field_relative_pose;
+}
+
+frc::Transform3d ExtrinsicsJsonToCameraToRobot(nlohmann::json extrinsics_json) {
+  frc::Transform3d robot_to_camera(
+      units::meter_t{extrinsics_json["translation_x"]},
+      units::meter_t{extrinsics_json["translation_y"]},
+      units::meter_t{extrinsics_json["translation_z"]},
+      frc::Rotation3d(units::radian_t{extrinsics_json["rotation_x"]},
+                      units::radian_t{extrinsics_json["rotation_y"]},
+                      units::radian_t{extrinsics_json["rotation_z"]}));
+  return robot_to_camera.Inverse();
+}
+
+std::vector<tag_detection_t> ToFeildRelitivePosition(
+    std::vector<tag_detection_t> detections, frc::Transform3d camera_to_robot,
+    frc::AprilTagFieldLayout apriltag_layout, bool verbose) {
+  for (size_t i = 0; i < detections.size(); ++i) {
+    detections[i] = ToFeildRelitivePosition(detections[i], camera_to_robot,
+                                            apriltag_layout, verbose);
+  }
+  return detections;
+}
+}  // namespace localization

src/localization/get_field_relitive_position.h

@@ -0,0 +1,25 @@
+#pragma once
+#include <frc/apriltag/AprilTagFieldLayout.h>
+#include <frc/geometry/Transform3d.h>
+#include <vector>
+#include "nlohmann/json_fwd.hpp"
+#include "src/localization/position.h"
+namespace localization {
+
+static frc::AprilTagFieldLayout kapriltag_layout =
+    frc::AprilTagFieldLayout::LoadField(
+        frc::AprilTagField::k2025ReefscapeAndyMark);
+
+std::vector<tag_detection_t> ToFeildRelitivePosition(
+    std::vector<tag_detection_t> detections, frc::Transform3d camera_to_robot,
+    frc::AprilTagFieldLayout apriltag_layout = kapriltag_layout,
+    bool verbose = false);
+
+tag_detection_t ToFeildRelitivePosition(
+    tag_detection_t tag_relative_position, frc::Transform3d camera_to_robot,
+    frc::AprilTagFieldLayout apriltag_layout = kapriltag_layout,
+    bool verbose = false);
+
+frc::Transform3d ExtrinsicsJsonToCameraToRobot(nlohmann::json extrinsics_json);
+
+}  // namespace localization

src/localization/gpu_apriltag_detector.cc

@@ -0,0 +1,148 @@
+#include "src/localization/gpu_apriltag_detector.h"
+#include <frc/geometry/Transform3d.h>
+#include <cmath>
+#include <opencv2/calib3d.hpp>
+#include <opencv2/imgproc.hpp>
+#include "apriltag/apriltag.h"
+#include "apriltag/tag36h11.h"
+#include "third_party/971apriltag/971apriltag.h"
+
+namespace localization {
+using json = nlohmann::json;
+
+constexpr double RadianToDegree(double radian) {
+  return radian * (180 / M_PI);
+}
+
+template <>
+frc971::apriltag::CameraMatrix
+camera_matrix_from_json<frc971::apriltag::CameraMatrix>(json intrinsics) {
+  frc971::apriltag::CameraMatrix camera_matrix = {.fx = intrinsics["fx"],
+                                                  .cx = intrinsics["cx"],
+                                                  .fy = intrinsics["fy"],
+                                                  .cy = intrinsics["cy"]};
+  return camera_matrix;
+}
+
+template <>
+frc971::apriltag::DistCoeffs
+distortion_coefficients_from_json<frc971::apriltag::DistCoeffs>(
+    json intrinsics) {
+  frc971::apriltag::DistCoeffs distortion_coefficients = {
+      .k1 = intrinsics["k1"],
+      .k2 = intrinsics["k2"],
+      .p1 = intrinsics["p1"],
+      .p2 = intrinsics["p2"],
+      .k3 = intrinsics["k3"]};
+
+  return distortion_coefficients;
+}
+
+template <>
+cv::Mat camera_matrix_from_json<cv::Mat>(json intrinsics) {
+  cv::Mat camera_matrix =
+      (cv::Mat_<double>(3, 3) << intrinsics["fx"], 0, intrinsics["cx"], 0,
+       intrinsics["fy"], intrinsics["cy"], 0, 0, 1);
+  return camera_matrix;
+}
+
+template <>
+cv::Mat distortion_coefficients_from_json<cv::Mat>(json intrinsics) {
+  cv::Mat distortion_coefficients =
+      (cv::Mat_<double>(1, 5) << intrinsics["k1"], intrinsics["k2"],
+       intrinsics["p1"], intrinsics["p2"], intrinsics["k3"]);
+  return distortion_coefficients;
+}
+
+GPUAprilTagDetector::GPUAprilTagDetector(
+    uint image_width, uint image_height, nlohmann::json intrinsics,
+    std::vector<cv::Point3f> apriltag_dimensions, bool verbose)
+    : camera_matrix_(camera_matrix_from_json<cv::Mat>(intrinsics)),
+      distortion_coefficients_(
+          distortion_coefficients_from_json<cv::Mat>(intrinsics)),
+      apriltag_dimensions_(apriltag_dimensions),
+      verbose_(verbose) {
+
+  apriltag_detector_ = apriltag_detector_create();
+
+  apriltag_detector_add_family_bits(apriltag_detector_, tag36h11_create(), 1);
+
+  apriltag_detector_->nthreads = 6;
+  apriltag_detector_->wp = workerpool_create(apriltag_detector_->nthreads);
+  apriltag_detector_->qtp.min_white_black_diff = 4;
+  apriltag_detector_->debug = false;
+
+  gpu_detector_ = std::make_unique<frc971::apriltag::GpuDetector>(
+      image_width, image_height, apriltag_detector_,
+      camera_matrix_from_json<frc971::apriltag::CameraMatrix>(intrinsics),
+      distortion_coefficients_from_json<frc971::apriltag::DistCoeffs>(
+          intrinsics));
+}
+std::vector<tag_detection_t> GPUAprilTagDetector::GetTagDetections(
+    camera::timestamped_frame_t& timestamped_frame) {
+  if (timestamped_frame.frame.channels() == 1) {
+    gpu_detector_->DetectGrayHost(
+        (unsigned char*)timestamped_frame.frame.ptr());
+  } else if (timestamped_frame.frame.channels() == 3) {
+    cv::Mat gray;
+    cv::cvtColor(timestamped_frame.frame, gray, cv::COLOR_BGR2GRAY);
+    gpu_detector_->DetectGrayHost((unsigned char*)gray.ptr());
+  }
+  const zarray_t* detections = gpu_detector_->Detections();
+  std::vector<tag_detection_t> estimates;
+
+  if (zarray_size(detections)) {
+    for (int i = 0; i < zarray_size(detections); ++i) {
+      apriltag_detection_t* gpu_detection;
+      zarray_get(detections, i, &gpu_detection);
+
+      std::vector<cv::Point2f> imagePoints;
+      for (int i = 0; i < 4; ++i) {
+
+        imagePoints.emplace_back(gpu_detection->p[i][0],
+                                 gpu_detection->p[i][1]);
+      }
+      cv::Mat rvec = cv::Mat::zeros(3, 1, CV_64FC1);  // output rotation vector
+      cv::Mat tvec =
+          cv::Mat::zeros(3, 1, CV_64FC1);  // output translation vector
+      cv::solvePnP(apriltag_dimensions_, imagePoints, camera_matrix_,
+                   distortion_coefficients_, rvec, tvec, false,
+                   cv::SOLVEPNP_IPPE_SQUARE);
+
+      tag_detection_t estimate;
+      // Currently we do not use transation z, rotation x and rotation y
+      // Converting to wpi coordinates
+      const double translation_x = tvec.ptr<double>()[2];
+      const double translation_y = tvec.ptr<double>()[0];
+      const double translation_z = tvec.ptr<double>()[1];
+
+      const double rotation_x = rvec.ptr<double>()[2];
+      const double rotation_y = rvec.ptr<double>()[0];
+      const double rotation_z = rvec.ptr<double>()[1];
+
+      estimate.pose =
+          frc::Pose3d(frc::Translation3d(units::meter_t{translation_x},
+                                         units::meter_t{translation_y},
+                                         units::meter_t{translation_z}),
+                      frc::Rotation3d(units::radian_t{rotation_x},
+                                      units::radian_t{rotation_y},
+                                      units::radian_t{rotation_z}));
+
+      estimate.distance = std::hypot(translation_x, translation_y);
+      estimate.tag_id = gpu_detection->id;
+
+      estimate.timestamp = timestamped_frame.timestamp;
+
+      estimates.push_back(estimate);
+    }
+  }
+  return estimates;
+}
+
+GPUAprilTagDetector::~GPUAprilTagDetector() {
+  if (apriltag_detector_ != nullptr) {
+    apriltag_detector_destroy(apriltag_detector_);
+  }
+}
+
+}  // namespace localization