Undistort points before running pose estimation

Signed-off-by: Maxwell Henderson <mxwhenderson@gmail.com>
Change-Id: Iaa332d8ad3518340ad93b39a2c558fad90cfae35
diff --git a/frc971/orin/apriltag.h b/frc971/orin/apriltag.h
index 3309996..8b5f0b2 100644
--- a/frc971/orin/apriltag.h
+++ b/frc971/orin/apriltag.h
@@ -193,6 +193,9 @@
     distortion_coefficients_ = distortion_coefficients;
   }
 
+  static void UnDistort(double *u, double *v, const CameraMatrix *camera_matrix,
+                        const DistCoeffs *distortion_coefficients);
+
  private:
   void UpdateFitQuads();
 
diff --git a/frc971/orin/apriltag_detect.cc b/frc971/orin/apriltag_detect.cc
index 9eeb9da..235a6d4 100644
--- a/frc971/orin/apriltag_detect.cc
+++ b/frc971/orin/apriltag_detect.cc
@@ -334,8 +334,9 @@
 
 // We're undistorting using math found from this github page
 // https://yangyushi.github.io/code/2020/03/04/opencv-undistort.html
-void UnDistort(double *x, double *y, CameraMatrix *camera_matrix,
-               DistCoeffs *distortion_coefficients) {
+void GpuDetector::UnDistort(double *u, double *v,
+                            const CameraMatrix *camera_matrix,
+                            const DistCoeffs *distortion_coefficients) {
   const double k1 = distortion_coefficients->k1;
   const double k2 = distortion_coefficients->k2;
   const double p1 = distortion_coefficients->p1;
@@ -347,8 +348,11 @@
   const double fy = camera_matrix->fy;
   const double cy = camera_matrix->cy;
 
-  double xP = (*x - cx) / fx;
-  double yP = (*y - cy) / fy;
+  const double xPP = (*u - cx) / fx;
+  const double yPP = (*v - cy) / fy;
+
+  double xP = xPP;
+  double yP = yPP;
 
   double x0 = xP;
   double y0 = yP;
@@ -391,8 +395,8 @@
             << " (" << prev_x << ", " << prev_y << ")";
   }
 
-  *x = xP * fx + cx;
-  *y = yP * fy + cy;
+  *u = xP * fx + cx;
+  *v = yP * fy + cy;
 }
 
 // Mostly stolen from aprilrobotics, but modified to implement the dewarp.
@@ -494,7 +498,8 @@
       double bestx = x0 + n0 * nx;
       double besty = y0 + n0 * ny;
 
-      UnDistort(&bestx, &besty, camera_matrix, distortion_coefficients);
+      GpuDetector::UnDistort(&bestx, &besty, camera_matrix,
+                             distortion_coefficients);
 
       // update our line fit statistics
       Mx += bestx;
diff --git a/frc971/orin/gpu_apriltag.cc b/frc971/orin/gpu_apriltag.cc
index 27cab4d..55c5cd5 100644
--- a/frc971/orin/gpu_apriltag.cc
+++ b/frc971/orin/gpu_apriltag.cc
@@ -74,8 +74,10 @@
       intrinsics_(frc971::vision::CameraIntrinsics(calibration_)),
       extrinsics_(frc971::vision::CameraExtrinsics(calibration_)),
       dist_coeffs_(frc971::vision::CameraDistCoeffs(calibration_)),
-      gpu_detector_(width, height, tag_detector_, GetCameraMatrix(calibration_),
-                    GetDistCoeffs(calibration_)),
+      distortion_camera_matrix_(GetCameraMatrix(calibration_)),
+      distortion_coefficients_(GetDistCoeffs(calibration_)),
+      gpu_detector_(width, height, tag_detector_, distortion_camera_matrix_,
+                    distortion_coefficients_),
       image_callback_(
           event_loop, channel_name,
           [this](cv::Mat image_color_mat,
@@ -136,23 +138,15 @@
 }
 
 void ApriltagDetector::UndistortDetection(apriltag_detection_t *det) const {
-  // 4 corners
-  constexpr size_t kRows = 4;
-  // 2d points
-  constexpr size_t kCols = 2;
-
-  cv::Mat distorted_points(kRows, kCols, CV_64F, det->p);
-  cv::Mat undistorted_points = cv::Mat::zeros(kRows, kCols, CV_64F);
-
-  // Undistort the april tag points
-  cv::undistortPoints(distorted_points, undistorted_points, intrinsics_,
-                      dist_coeffs_, cv::noArray(), projection_matrix_);
-
   // Copy the undistorted points into det
-  for (size_t i = 0; i < kRows; i++) {
-    for (size_t j = 0; j < kCols; j++) {
-      det->p[i][j] = undistorted_points.at<double>(i, j);
-    }
+  for (size_t i = 0; i < 4; i++) {
+    double u = det->p[i][0];
+    double v = det->p[i][1];
+
+    GpuDetector::UnDistort(&u, &v, &distortion_camera_matrix_,
+                           &distortion_coefficients_);
+    det->p[i][0] = u;
+    det->p[i][1] = v;
   }
 }
 
@@ -267,7 +261,6 @@
       // First create an apriltag_detection_info_t struct using your known
       // parameters.
       apriltag_detection_info_t info;
-      info.det = gpu_detection;
       info.tagsize = 6.5 * 0.0254;
 
       info.fx = intrinsics_.at<double>(0, 0);
@@ -284,6 +277,10 @@
 
       UndistortDetection(gpu_detection);
 
+      // We're setting this here to use the undistorted corner points in pose
+      // estimation.
+      info.det = gpu_detection;
+
       const aos::monotonic_clock::time_point before_pose_estimation =
           aos::monotonic_clock::now();
 
diff --git a/frc971/orin/gpu_apriltag.h b/frc971/orin/gpu_apriltag.h
index 39d2b60..5569e5a 100644
--- a/frc971/orin/gpu_apriltag.h
+++ b/frc971/orin/gpu_apriltag.h
@@ -81,6 +81,11 @@
   std::optional<cv::Mat> extrinsics_;
   cv::Mat dist_coeffs_;
 
+  // The distortion constants passed into gpu_detector_ and used later on to do
+  // undistort before passing corners into apriltags pose estimation.
+  CameraMatrix distortion_camera_matrix_;
+  DistCoeffs distortion_coefficients_;
+
   frc971::apriltag::GpuDetector gpu_detector_;
   cv::Size image_size_;