Dewarp refined tag edges

Implements undistort on our bestx and besty

Signed-off-by: Maxwell Henderson <mxwhenderson@gmail.com>
Change-Id: I3594b47161cb171de1597a66e21e5d28404109da
diff --git a/WORKSPACE b/WORKSPACE
index 7809594..b00de63 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1577,6 +1577,20 @@
     url = "https://software.frc971.org/Build-Dependencies/orin_large_gs_apriltag.bfbs",
 )
 
+http_file(
+    name = "orin_capture_24_04",
+    downloaded_file_path = "orin_capture_24_04.bfbs",
+    sha256 = "719edb1d1394c13c1b55d02cf35c277e1d4c2111f4eb4220b28addc08634488a",
+    url = "https://software.frc971.org/Build-Dependencies/orin-capture-24-04-2024.02.14.bfbs",
+)
+
+http_file(
+    name = "orin_capture_24_04_side",
+    downloaded_file_path = "orin_capture_24_04_side.bfbs",
+    sha256 = "4747cc98f8794d6570cb12a3171d7984e358581914a28b43fb6bb8b9bd7a10ac",
+    url = "https://software.frc971.org/Build-Dependencies/orin-capture-24-04-side-2024.02.17.bfbs",
+)
+
 http_archive(
     name = "libedgetpu",
     build_file = "//third_party:libedgetpu/libedgetpu.BUILD",
diff --git a/frc971/orin/BUILD b/frc971/orin/BUILD
index 6383423..06b1d6c 100644
--- a/frc971/orin/BUILD
+++ b/frc971/orin/BUILD
@@ -61,6 +61,8 @@
     ],
     data = [
         "@apriltag_test_bfbs_images",
+        "@orin_capture_24_04//file",
+        "@orin_capture_24_04_side//file",
         "@orin_image_apriltag//file",
         "@orin_large_image_apriltag//file",
     ],
diff --git a/frc971/orin/apriltag.cc b/frc971/orin/apriltag.cc
index dd253d1..92513c2 100644
--- a/frc971/orin/apriltag.cc
+++ b/frc971/orin/apriltag.cc
@@ -113,7 +113,9 @@
 }  // namespace
 
 GpuDetector::GpuDetector(size_t width, size_t height,
-                         apriltag_detector_t *tag_detector)
+                         apriltag_detector_t *tag_detector,
+                         CameraMatrix camera_matrix,
+                         DistCoeffs distortion_coefficients)
     : width_(width),
       height_(height),
       tag_detector_(tag_detector),
@@ -141,6 +143,8 @@
       compressed_peaks_device_(line_fit_points_device_.size()),
       sorted_compressed_peaks_device_(line_fit_points_device_.size()),
       peak_extents_device_(kMaxBlobs),
+      camera_matrix_(camera_matrix),
+      distortion_coefficients_(distortion_coefficients),
       fit_quads_device_(kMaxBlobs),
       radix_sort_tmpstorage_device_(RadixSortScratchSpace<QuadBoundaryPoint>(
           sorted_union_marker_pair_device_.size())),
diff --git a/frc971/orin/apriltag.h b/frc971/orin/apriltag.h
index d3e041b..3309996 100644
--- a/frc971/orin/apriltag.h
+++ b/frc971/orin/apriltag.h
@@ -59,6 +59,21 @@
   uint32_t blob_index;
 };
 
+struct CameraMatrix {
+  double fx;
+  double cx;
+  double fy;
+  double cy;
+};
+
+struct DistCoeffs {
+  double k1;
+  double k2;
+  double p1;
+  double p2;
+  double k3;
+};
+
 // GPU based april tag detector.
 class GpuDetector {
  public:
@@ -67,8 +82,8 @@
 
   // Constructs a detector, reserving space for detecting tags of the provided
   // with and height, using the provided detector options.
-  GpuDetector(size_t width, size_t height, apriltag_detector_t *tag_detector);
-
+  GpuDetector(size_t width, size_t height, apriltag_detector_t *tag_detector,
+              CameraMatrix camera_matrix, DistCoeffs distortion_coefficients);
   virtual ~GpuDetector();
 
   // Detects april tags in the provided image.
@@ -168,6 +183,16 @@
 
   void AdjustCenter(float corners[4][2]) const;
 
+  // TODO(max): We probably don't want to use these after our test images are
+  // just orin images
+  void SetCameraMatrix(CameraMatrix camera_matrix) {
+    camera_matrix_ = camera_matrix;
+  }
+
+  void SetDistortionCoefficients(DistCoeffs distortion_coefficients) {
+    distortion_coefficients_ = distortion_coefficients;
+  }
+
  private:
   void UpdateFitQuads();
 
@@ -292,6 +317,9 @@
   GpuMemory<int> num_quad_peaked_quads_device_{/* allocate 1 integer...*/ 1};
   GpuMemory<PeakExtents> peak_extents_device_;
 
+  CameraMatrix camera_matrix_;
+  DistCoeffs distortion_coefficients_;
+
   GpuMemory<FitQuad> fit_quads_device_;
 
   std::vector<FitQuad> fit_quads_host_;
diff --git a/frc971/orin/apriltag_detect.cc b/frc971/orin/apriltag_detect.cc
index 69c93cc..5671a0d 100644
--- a/frc971/orin/apriltag_detect.cc
+++ b/frc971/orin/apriltag_detect.cc
@@ -295,11 +295,79 @@
   zarray_t *detections;
 
   image_u8_t *im_samples;
+
+  CameraMatrix *camera_matrix;
+  DistCoeffs *distortion_coefficients;
 };
 
+// Dewarps points from the image based on various constants
+// Algorithm mainly taken from
+// https://docs.opencv.org/4.x/dc/dbb/tutorial_py_calibration.html
+void ReDistort(double *x, double *y, CameraMatrix *camera_matrix,
+               DistCoeffs *distortion_coefficients) {
+  double k1 = distortion_coefficients->k1;
+  double k2 = distortion_coefficients->k2;
+  double p1 = distortion_coefficients->p1;
+  double p2 = distortion_coefficients->p2;
+  double k3 = distortion_coefficients->k3;
+
+  double fx = camera_matrix->fx;
+  double cx = camera_matrix->cx;
+  double fy = camera_matrix->fy;
+  double cy = camera_matrix->cy;
+
+  double xP = (*x - cx) / fx;
+  double yP = (*y - cy) / fy;
+
+  double rSq = xP * xP + yP * yP;
+
+  double linCoef = 1 + k1 * rSq + k2 * rSq * rSq + k3 * rSq * rSq * rSq;
+  double xPP = xP * linCoef + 2 * p1 * xP * yP + p2 * (rSq + 2 * xP * xP);
+  double yPP = yP * linCoef + p1 * (rSq + 2 * yP * yP) + 2 * p2 * xP * yP;
+
+  *x = xPP * fx + cx;
+  *y = yPP * fy + cy;
+}
+
+// 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) {
+  double k1 = distortion_coefficients->k1;
+  double k2 = distortion_coefficients->k2;
+  double p1 = distortion_coefficients->p1;
+  double p2 = distortion_coefficients->p2;
+  double k3 = distortion_coefficients->k3;
+
+  double fx = camera_matrix->fx;
+  double cx = camera_matrix->cx;
+  double fy = camera_matrix->fy;
+  double cy = camera_matrix->cy;
+
+  double xP = (*x - cx) / fx;
+  double yP = (*y - cy) / fy;
+
+  double x0 = xP;
+  double y0 = yP;
+
+  for (int i = 0; i < 3; i++) {
+    double rSq = xP * xP + yP * yP;
+    double linCoef = 1 + k1 * rSq + k2 * rSq * rSq + k3 * rSq * rSq * rSq;
+    double kInv = 1 / linCoef;
+    double dx = 2 * p1 * xP * yP + p2 * (rSq + k3 * rSq * rSq * rSq);
+    double dy = p1 * (rSq + 2 * yP * yP) + 2 * p2 * xP * yP;
+    xP = (x0 - dx) * kInv;
+    yP = (y0 - dy) * kInv;
+  }
+
+  *x = xP * fx + cx;
+  *y = yP * fy + cy;
+}
+
 // Mostly stolen from aprilrobotics, but modified to implement the dewarp.
 void RefineEdges(apriltag_detector_t *td, image_u8_t *im_orig,
-                 struct quad *quad) {
+                 struct quad *quad, CameraMatrix *camera_matrix,
+                 DistCoeffs *distortion_coefficients) {
   double lines[4][4];  // for each line, [Ex Ey nx ny]
 
   for (int edge = 0; edge < 4; edge++) {
@@ -395,6 +463,8 @@
       double bestx = x0 + n0 * nx;
       double besty = y0 + n0 * ny;
 
+      UnDistort(&bestx, &besty, camera_matrix, distortion_coefficients);
+
       // update our line fit statistics
       Mx += bestx;
       My += besty;
@@ -440,8 +510,12 @@
       // Compute intersection. Note that line i represents the line from corner
       // i to (i+1)&3, so the intersection of line i with line (i+1)&3
       // represents corner (i+1)&3.
-      quad->p[(i + 1) & 3][0] = lines[i][0] + L0 * A00;
-      quad->p[(i + 1) & 3][1] = lines[i][1] + L0 * A10;
+      double px = lines[i][0] + L0 * A00;
+      double py = lines[i][1] + L0 * A10;
+
+      ReDistort(&px, &py, camera_matrix, distortion_coefficients);
+      quad->p[(i + 1) & 3][0] = px;
+      quad->p[(i + 1) & 3][1] = py;
     } else {
       // this is a bad sign. We'll just keep the corner we had.
       //            debug_print("bad det: %15f %15f %15f %15f %15f\n", A00, A11,
@@ -465,7 +539,36 @@
     quad_original.Hinv = nullptr;
 
     if (td->refine_edges) {
-      RefineEdges(td, im, &quad_original);
+      RefineEdges(td, im, &quad_original, task->camera_matrix,
+                  task->distortion_coefficients);
+    }
+
+    if (td->debug) {
+      image_u8_t *im_quads = image_u8_copy(im);
+      image_u8_darken(im_quads);
+      image_u8_darken(im_quads);
+
+      srandom(0);
+
+      const int bias = 100;
+      int color = bias + (random() % (255 - bias));
+
+      image_u8_draw_line(im_quads, quad_original.p[0][0], quad_original.p[0][1],
+                         quad_original.p[1][0], quad_original.p[1][1], color,
+                         1);
+      image_u8_draw_line(im_quads, quad_original.p[1][0], quad_original.p[1][1],
+                         quad_original.p[2][0], quad_original.p[2][1], color,
+                         1);
+      image_u8_draw_line(im_quads, quad_original.p[2][0], quad_original.p[2][1],
+                         quad_original.p[3][0], quad_original.p[3][1], color,
+                         1);
+      image_u8_draw_line(im_quads, quad_original.p[3][0], quad_original.p[3][1],
+                         quad_original.p[0][0], quad_original.p[0][1], color,
+                         1);
+
+      image_u8_write_pnm(
+          im_quads,
+          std::string("/tmp/quad" + std::to_string(quadidx) + ".pnm").c_str());
     }
 
     quad_decode_index(td, &quad_original, im, task->im_samples,
@@ -504,6 +607,8 @@
     tasks[ntasks].td = tag_detector_;
     tasks[ntasks].im = &im_orig;
     tasks[ntasks].detections = detections_;
+    tasks[ntasks].camera_matrix = &camera_matrix_;
+    tasks[ntasks].distortion_coefficients = &distortion_coefficients_;
 
     tasks[ntasks].im_samples = nullptr;
 
diff --git a/frc971/orin/cuda_april_tag_test.cc b/frc971/orin/cuda_april_tag_test.cc
index b30943f..638bf6a 100644
--- a/frc971/orin/cuda_april_tag_test.cc
+++ b/frc971/orin/cuda_april_tag_test.cc
@@ -8,6 +8,8 @@
 #include "third_party/apriltag/apriltag.h"
 #include "third_party/apriltag/common/unionfind.h"
 #include "third_party/apriltag/tag16h5.h"
+#include "third_party/apriltag/tag36h11.h"
+#include <opencv2/calib3d.hpp>
 #include <opencv2/highgui.hpp>
 
 #include "aos/flatbuffer_merge.h"
@@ -243,17 +245,35 @@
   return tag_detector;
 }
 
+// TODO(max): Create a function which will take in the calibration data
+CameraMatrix create_camera_matrix() {
+  return CameraMatrix{
+      1,
+      1,
+      1,
+      1,
+  };
+}
+
+DistCoeffs create_distortion_coefficients() {
+  return DistCoeffs{
+      0, 0, 0, 0, 0,
+  };
+}
+
 class CudaAprilTagDetector {
  public:
-  CudaAprilTagDetector(size_t width, size_t height)
-      : tag_family_(tag16h5_create()),
+  CudaAprilTagDetector(size_t width, size_t height,
+                       apriltag_family_t *tag_family = tag16h5_create())
+      : tag_family_(tag_family),
         tag_detector_(MakeTagDetector(tag_family_)),
         gray_cuda_(cv::Size(width, height), CV_8UC1),
         decimated_cuda_(gray_cuda_.size() / 2, CV_8UC1),
         thresholded_cuda_(decimated_cuda_.size(), CV_8UC1),
         union_markers_(decimated_cuda_.size(), CV_32SC1),
         union_markers_size_(decimated_cuda_.size(), CV_32SC1),
-        gpu_detector_(width, height, tag_detector_),
+        gpu_detector_(width, height, tag_detector_, create_camera_matrix(),
+                      create_distortion_coefficients()),
         width_(width),
         height_(height) {
     // Report out info about our GPU.
@@ -1716,6 +1736,8 @@
     }
   }
 
+  void set_undistort(bool value) { undistort_ = value; }
+
   void CheckDetections(zarray_t *aprilrobotics_detections,
                        const zarray_t *_gpu_detections) {
     zarray_t *gpu_detections = zarray_copy(_gpu_detections);
@@ -1757,7 +1779,7 @@
 
       // TODO(austin): Crank down the thresholds and figure out why these
       // deviate.  It should be the same function for both at this point.
-      const double threshold = valid ? 2e-3 : 1e-1;
+      const double threshold = undistort_ ? 15.0 : (valid ? 2e-3 : 1e-1);
 
       CHECK_EQ(aprilrobotics_detection->id, gpu_detection->id);
       CHECK_EQ(aprilrobotics_detection->hamming, gpu_detection->hamming);
@@ -1869,7 +1891,9 @@
     zarray_t *april_clusters =
         OrderAprilroboticsLikeCuda(thresholded_im, uf, cuda_grouped_points);
 
-    CheckQuads(april_clusters, sorted_selected_blobs_cuda_, fit_quads_);
+    if (!undistort_) {
+      CheckQuads(april_clusters, sorted_selected_blobs_cuda_, fit_quads_);
+    }
 
     const zarray_t *gpu_detections = gpu_detector_.Detections();
     CheckDetections(aprilrobotics_detections_, gpu_detections);
@@ -2045,6 +2069,14 @@
     }
   }
 
+  // Sets the camera constants for camera 24-04
+  void SetCameraFourConstants() {
+    gpu_detector_.SetCameraMatrix(
+        CameraMatrix{642.80365, 718.017517, 642.83667, 555.022461});
+    gpu_detector_.SetDistortionCoefficients(
+        DistCoeffs{-0.239969, 0.055889, 0.000086, 0.000099, -0.005468});
+  }
+
  private:
   apriltag_family_t *tag_family_;
   apriltag_detector_t *tag_detector_;
@@ -2082,6 +2114,8 @@
 
   bool normal_border_ = false;
   bool reversed_border_ = false;
+
+  bool undistort_ = false;
   int min_tag_width_ = 1000000;
 };
 
@@ -2236,4 +2270,52 @@
   EXPECT_EQ(overall_count, MaxRankedIndex());
 }
 
+// Tests our Undistort is working properly
+TEST_F(AprilDetectionTest, Undistort) {
+  auto image = ReadImage("orin_capture_24_04/file/orin_capture_24_04.bfbs");
+
+  LOG(INFO) << "Image is: " << image.message().cols() << " x "
+            << image.message().rows();
+
+  CudaAprilTagDetector cuda_detector(image.message().cols(),
+                                     image.message().rows(), tag36h11_create());
+
+  cuda_detector.set_undistort(true);
+
+  const cv::Mat color_image = ToMat(&image.message());
+
+  cuda_detector.SetCameraFourConstants();
+
+  cuda_detector.DetectGPU(color_image.clone());
+  cuda_detector.DetectCPU(color_image.clone());
+  cuda_detector.Check(color_image.clone());
+  if (FLAGS_debug) {
+    cuda_detector.WriteDebug(color_image);
+  }
+}
+
+// Tests our Undistort is working properly with a tag at the edge of the image
+TEST_F(AprilDetectionTest, UndistortEdge) {
+  auto image =
+      ReadImage("orin_capture_24_04_side/file/orin_capture_24_04_side.bfbs");
+
+  LOG(INFO) << "Image is: " << image.message().cols() << " x "
+            << image.message().rows();
+
+  CudaAprilTagDetector cuda_detector(image.message().cols(),
+                                     image.message().rows(), tag36h11_create());
+
+  cuda_detector.set_undistort(true);
+
+  const cv::Mat color_image = ToMat(&image.message());
+
+  cuda_detector.SetCameraFourConstants();
+
+  cuda_detector.DetectGPU(color_image.clone());
+  cuda_detector.DetectCPU(color_image.clone());
+  cuda_detector.Check(color_image.clone());
+  if (FLAGS_debug) {
+    cuda_detector.WriteDebug(color_image);
+  }
+}
 }  // namespace frc971::apriltag::testing
diff --git a/frc971/orin/gpu_apriltag.cc b/frc971/orin/gpu_apriltag.cc
index 89625b6..2857de7 100644
--- a/frc971/orin/gpu_apriltag.cc
+++ b/frc971/orin/gpu_apriltag.cc
@@ -40,15 +40,42 @@
 
 namespace chrono = std::chrono;
 
+CameraMatrix GetCameraMatrix(
+    const frc971::vision::calibration::CameraCalibration *calibration) {
+  auto intrinsics = calibration->intrinsics();
+  return CameraMatrix{
+      .fx = intrinsics->Get(0),
+      .cx = intrinsics->Get(2),
+      .fy = intrinsics->Get(4),
+      .cy = intrinsics->Get(5),
+  };
+}
+
+DistCoeffs GetDistCoeffs(
+    const frc971::vision::calibration::CameraCalibration *calibration) {
+  auto dist_coeffs = calibration->dist_coeffs();
+  return DistCoeffs{
+      .k1 = dist_coeffs->Get(0),
+      .k2 = dist_coeffs->Get(1),
+      .p1 = dist_coeffs->Get(2),
+      .p2 = dist_coeffs->Get(3),
+      .k3 = dist_coeffs->Get(4),
+  };
+}
+
 ApriltagDetector::ApriltagDetector(
     aos::EventLoop *event_loop, std::string_view channel_name,
     const frc971::vision::calibration::CameraCalibration *calibration,
     size_t width, size_t height)
     : tag_family_(tag36h11_create()),
       tag_detector_(MakeTagDetector(tag_family_)),
-      gpu_detector_(width, height, tag_detector_),
       node_name_(event_loop->node()->name()->string_view()),
       calibration_(calibration),
+      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_)),
       image_callback_(
           event_loop, channel_name,
           [this](cv::Mat image_color_mat,
@@ -63,10 +90,6 @@
       rejections_(0) {
   image_callback_.set_format(frc971::vision::ImageCallback::Format::YUYV2);
 
-  extrinsics_ = frc971::vision::CameraExtrinsics(calibration_);
-  intrinsics_ = frc971::vision::CameraIntrinsics(calibration_);
-  dist_coeffs_ = frc971::vision::CameraDistCoeffs(calibration_);
-
   projection_matrix_ = cv::Mat::zeros(3, 4, CV_64F);
   intrinsics_.rowRange(0, 3).colRange(0, 3).copyTo(
       projection_matrix_.rowRange(0, 3).colRange(0, 3));
diff --git a/frc971/orin/gpu_apriltag.h b/frc971/orin/gpu_apriltag.h
index 09c8579..39d2b60 100644
--- a/frc971/orin/gpu_apriltag.h
+++ b/frc971/orin/gpu_apriltag.h
@@ -73,7 +73,6 @@
  private:
   apriltag_family_t *tag_family_;
   apriltag_detector_t *tag_detector_;
-  frc971::apriltag::GpuDetector gpu_detector_;
   std::string_view node_name_;
 
   const frc971::vision::calibration::CameraCalibration *calibration_;
@@ -81,6 +80,8 @@
   cv::Mat projection_matrix_;
   std::optional<cv::Mat> extrinsics_;
   cv::Mat dist_coeffs_;
+
+  frc971::apriltag::GpuDetector gpu_detector_;
   cv::Size image_size_;
 
   frc971::vision::ImageCallback image_callback_;