Add aprilrobotics detector app

Change-Id: I2446710755fd5226c3e1f62fa2a4d0e8c7e27c17
Signed-off-by: Maxwell Henderson <mxwhenderson@gmail.com>
diff --git a/frc971/vision/target_map.fbs b/frc971/vision/target_map.fbs
index 38bab6d..37260a6 100644
--- a/frc971/vision/target_map.fbs
+++ b/frc971/vision/target_map.fbs
@@ -26,4 +26,6 @@
 
   // Unique name of the field
   field_name:string (id: 1);
-}
\ No newline at end of file
+}
+
+root_type TargetMap;
diff --git a/third_party/apriltag/BUILD b/third_party/apriltag/BUILD
new file mode 100644
index 0000000..e4acc96
--- /dev/null
+++ b/third_party/apriltag/BUILD
@@ -0,0 +1,86 @@
+cc_library(
+    name = "apriltag",
+    srcs = [
+        "apriltag.c",
+        "apriltag_pose.c",
+        "apriltag_quad_thresh.c",
+        "common/g2d.c",
+        "common/getopt.c",
+        "common/homography.c",
+        "common/image_u8.c",
+        "common/image_u8x3.c",
+        "common/image_u8x4.c",
+        "common/matd.c",
+        "common/pam.c",
+        "common/pjpeg.c",
+        "common/pjpeg-idct.c",
+        "common/pnm.c",
+        "common/string_util.c",
+        "common/svd22.c",
+        "common/time_util.c",
+        "common/unionfind.c",
+        "common/workerpool.c",
+        "common/zarray.c",
+        "common/zhash.c",
+        "common/zmaxheap.c",
+        "tag16h5.c",
+        "tag25h9.c",
+        "tag36h11.c",
+        "tagCircle21h7.c",
+        "tagCircle49h12.c",
+        "tagCustom48h12.c",
+        "tagStandard41h12.c",
+        "tagStandard52h13.c",
+    ],
+    hdrs = [
+        "apriltag.h",
+        "apriltag_math.h",
+        "apriltag_pose.h",
+        "common/debug_print.h",
+        "common/g2d.h",
+        "common/getopt.h",
+        "common/homography.h",
+        "common/image_types.h",
+        "common/image_u8.h",
+        "common/image_u8x3.h",
+        "common/image_u8x4.h",
+        "common/matd.h",
+        "common/math_util.h",
+        "common/pam.h",
+        "common/pjpeg.h",
+        "common/pnm.h",
+        "common/postscript_utils.h",
+        "common/pthreads_cross.h",
+        "common/string_util.h",
+        "common/svd22.h",
+        "common/time_util.h",
+        "common/timeprofile.h",
+        "common/unionfind.h",
+        "common/workerpool.h",
+        "common/zarray.h",
+        "common/zhash.h",
+        "common/zmaxheap.h",
+        "tag16h5.h",
+        "tag25h9.h",
+        "tag36h11.h",
+        "tagCircle21h7.h",
+        "tagCircle49h12.h",
+        "tagCustom48h12.h",
+        "tagStandard41h12.h",
+        "tagStandard52h13.h",
+    ],
+    copts = [
+        "-Wno-unused-parameter",
+        "-Wno-cast-qual",
+        "-Wno-sign-compare",
+        "-Wno-unused-variable",
+        "-Wno-cast-align",
+        "-Wno-incompatible-pointer-types-discards-qualifiers",
+        "-Wno-format-nonliteral",
+    ],
+    includes = ["."],
+    target_compatible_with = [
+        "@platforms//cpu:arm64",
+    ],
+    visibility = ["//visibility:public"],
+)
diff --git a/third_party/apriltag/common/workerpool.c b/third_party/apriltag/common/workerpool.c
index a0170ef..29eccfc 100644
--- a/third_party/apriltag/common/workerpool.c
+++ b/third_party/apriltag/common/workerpool.c
@@ -41,20 +41,6 @@
 #include "workerpool.h"
 #include "debug_print.h"
 
-struct workerpool {
-    int nthreads;
-    zarray_t *tasks;
-    int taskspos;
-
-    pthread_t *threads;
-    int *status;
-
-    pthread_mutex_t mutex;
-    pthread_cond_t startcond;   // used to signal the availability of work
-    pthread_cond_t endcond;     // used to signal completion of all work
-
-    int end_count; // how many threads are done?
-};
 
 struct task
 {
diff --git a/third_party/apriltag/common/workerpool.h b/third_party/apriltag/common/workerpool.h
index 2c32ab1..a233b5b 100644
--- a/third_party/apriltag/common/workerpool.h
+++ b/third_party/apriltag/common/workerpool.h
@@ -31,6 +31,21 @@
 
 typedef struct workerpool workerpool_t;
 
+struct workerpool {
+  int nthreads;
+  zarray_t *tasks;
+  int taskspos;
+
+  pthread_t *threads;
+  int *status;
+
+  pthread_mutex_t mutex;
+  pthread_cond_t startcond;  // used to signal the availability of work
+  pthread_cond_t endcond;    // used to signal completion of all work
+
+  int end_count;  // how many threads are done?
+};
+
 // as a special case, if nthreads==1, no additional threads are
 // created, and workerpool_run will run synchronously.
 workerpool_t *workerpool_create(int nthreads);
@@ -41,7 +56,8 @@
 // runs all added tasks, waits for them to complete.
 void workerpool_run(workerpool_t *wp);
 
-// same as workerpool_run, except always single threaded. (mostly for debugging).
+// same as workerpool_run, except always single threaded. (mostly for
+// debugging).
 void workerpool_run_single(workerpool_t *wp);
 
 int workerpool_get_nthreads(workerpool_t *wp);
diff --git a/y2023/BUILD b/y2023/BUILD
index 0879de7..60d218b 100644
--- a/y2023/BUILD
+++ b/y2023/BUILD
@@ -7,7 +7,7 @@
     binaries = [
         "//y2020/vision:calibration",
         "//y2023/vision:viewer",
-        "//y2022/localizer:imu_main",
+        "//y2023/vision:aprilrobotics",
         "//y2022/localizer:localizer_main",
         "//aos/network:web_proxy_main",
         "//aos/events/logging:log_cat",
@@ -48,6 +48,7 @@
         "//aos/network:timestamp_fbs",
         "//frc971/input:robot_state_fbs",
         "//frc971/vision:vision_fbs",
+        "//frc971/vision:target_map_fbs",
     ],
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//visibility:public"],
@@ -73,6 +74,7 @@
             "//aos/network:remote_message_fbs",
             "//y2022/localizer:localizer_output_fbs",
             "//frc971/vision:calibration_fbs",
+            "//frc971/vision:target_map_fbs",
             "//frc971/vision:vision_fbs",
         ],
         target_compatible_with = ["@platforms//os:linux"],
@@ -102,6 +104,7 @@
         "//y2022/localizer:localizer_status_fbs",
         "//y2022/localizer:localizer_output_fbs",
         "//y2022/localizer:localizer_visualization_fbs",
+        "//frc971/vision:target_map_fbs",
     ],
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//visibility:public"],
@@ -121,6 +124,7 @@
         "//aos/network:remote_message_fbs",
         "//frc971/vision:calibration_fbs",
         "//frc971/vision:vision_fbs",
+        "//frc971/vision:target_map_fbs",
     ],
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//visibility:public"],
diff --git a/y2023/vision/BUILD b/y2023/vision/BUILD
index 6ef0657..25ae2b9 100644
--- a/y2023/vision/BUILD
+++ b/y2023/vision/BUILD
@@ -97,3 +97,25 @@
         "//third_party:opencv",
     ],
 )
+
+cc_binary(
+    name = "aprilrobotics",
+    srcs = [
+        "aprilrobotics.cc",
+        "aprilrobotics.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//y2023:__subpackages__"],
+    deps = [
+        ":calibration_data",
+        "//aos:init",
+        "//aos/events:shm_event_loop",
+        "//frc971/vision:calibration_fbs",
+        "//frc971/vision:charuco_lib",
+        "//frc971/vision:target_map_fbs",
+        "//frc971/vision:target_mapper",
+        "//frc971/vision:vision_fbs",
+        "//third_party:opencv",
+        "//third_party/apriltag",
+    ],
+)
diff --git a/y2023/vision/aprilrobotics.cc b/y2023/vision/aprilrobotics.cc
new file mode 100644
index 0000000..f8a8e38
--- /dev/null
+++ b/y2023/vision/aprilrobotics.cc
@@ -0,0 +1,197 @@
+#include "y2023/vision/aprilrobotics.h"
+
+DEFINE_string(config, "aos_config.json", "Path to the config file to use.");
+
+DEFINE_bool(
+    debug, false,
+    "If true, dump a ton of debug and crash on the first valid detection.");
+
+DEFINE_int32(team_number, 971,
+             "Use the calibration for a node with this team number");
+namespace y2023 {
+namespace vision {
+
+AprilRoboticsDetector::AprilRoboticsDetector(aos::EventLoop *event_loop,
+                                             std::string_view channel_name)
+    : calibration_data_(CalibrationData()),
+      ftrace_(),
+      image_callback_(event_loop, channel_name,
+                      [&](cv::Mat image_color_mat,
+                          const aos::monotonic_clock::time_point /*eof*/) {
+                        HandleImage(image_color_mat);
+                      }),
+      target_map_sender_(
+          event_loop->MakeSender<frc971::vision::TargetMap>("/camera")) {
+  tag_family_ = tag16h5_create();
+  tag_detector_ = apriltag_detector_create();
+
+  apriltag_detector_add_family_bits(tag_detector_, tag_family_, 1);
+  tag_detector_->nthreads = 6;
+  tag_detector_->wp = workerpool_create(tag_detector_->nthreads);
+  tag_detector_->qtp.min_white_black_diff = 5;
+  tag_detector_->debug = FLAGS_debug;
+
+  std::string hostname = aos::network::GetHostname();
+
+  // Check team string is valid
+  std::optional<uint16_t> pi_number = aos::network::ParsePiNumber(hostname);
+  std::optional<uint16_t> team_number =
+      aos::network::team_number_internal::ParsePiTeamNumber(hostname);
+  CHECK(pi_number) << "Unable to parse pi number from '" << hostname << "'";
+  CHECK(team_number);
+
+  calibration_ = FindCameraCalibration(&calibration_data_.message(),
+                                       "pi" + std::to_string(*pi_number));
+  intrinsics_ = CameraIntrinsics(calibration_);
+  camera_distortion_coeffs_ = CameraDistCoeffs(calibration_);
+
+  image_callback_.set_format(frc971::vision::ImageCallback::Format::GRAYSCALE);
+}
+
+AprilRoboticsDetector::~AprilRoboticsDetector() {
+  apriltag_detector_destroy(tag_detector_);
+  free(tag_family_);
+}
+
+void AprilRoboticsDetector::SetWorkerpoolAffinities() {
+  for (int i = 0; i < tag_detector_->wp->nthreads; i++) {
+    cpu_set_t affinity;
+    CPU_ZERO(&affinity);
+    CPU_SET(i, &affinity);
+    pthread_setaffinity_np(tag_detector_->wp->threads[i], sizeof(affinity),
+                           &affinity);
+    struct sched_param param;
+    param.sched_priority = 20;
+    int res = pthread_setschedparam(tag_detector_->wp->threads[i], SCHED_FIFO,
+                                    &param);
+    PCHECK(res == 0) << "Failed to set priority of threadpool threads";
+  }
+}
+
+void AprilRoboticsDetector::HandleImage(cv::Mat image_color_mat) {
+  std::vector<std::pair<apriltag_detection_t, apriltag_pose_t>> detections =
+      DetectTags(image_color_mat);
+
+  auto builder = target_map_sender_.MakeBuilder();
+  std::vector<flatbuffers::Offset<frc971::vision::TargetPoseFbs>> target_poses;
+  for (const auto &[detection, pose] : detections) {
+    target_poses.emplace_back(
+        BuildTargetPose(pose, detection.id, builder.fbb()));
+  }
+  const auto target_poses_offset = builder.fbb()->CreateVector(target_poses);
+  auto target_map_builder = builder.MakeBuilder<frc971::vision::TargetMap>();
+
+  target_map_builder.add_target_poses(target_poses_offset);
+  builder.CheckOk(builder.Send(target_map_builder.Finish()));
+}
+
+flatbuffers::Offset<frc971::vision::TargetPoseFbs>
+AprilRoboticsDetector::BuildTargetPose(
+    const apriltag_pose_t &pose,
+    frc971::vision::TargetMapper::TargetId target_id,
+    flatbuffers::FlatBufferBuilder *fbb) {
+  const auto T =
+      Eigen::Translation3d(pose.t->data[0], pose.t->data[1], pose.t->data[2]);
+  const auto rpy = frc971::vision::PoseUtils::RotationMatrixToEulerAngles(
+      Eigen::Matrix3d(pose.R->data));
+  return frc971::vision::CreateTargetPoseFbs(*fbb, target_id, T.x(), T.y(),
+                                             T.z(), rpy(0), rpy(1), rpy(2));
+}
+
+std::vector<std::pair<apriltag_detection_t, apriltag_pose_t>>
+AprilRoboticsDetector::DetectTags(cv::Mat image) {
+  const aos::monotonic_clock::time_point start_time =
+      aos::monotonic_clock::now();
+
+  image_u8_t im = {
+    .width = image.cols,
+    .height = image.rows,
+    .stride = image.cols,
+    .buf = image.data,
+  };
+
+  ftrace_.FormatMessage("Starting detect\n");
+  zarray_t *detections = apriltag_detector_detect(tag_detector_, &im);
+  ftrace_.FormatMessage("Done detecting\n");
+
+  std::vector<std::pair<apriltag_detection_t, apriltag_pose_t>> results;
+
+  for (int i = 0; i < zarray_size(detections); i++) {
+    apriltag_detection_t *det;
+    zarray_get(detections, i, &det);
+
+    if (det->decision_margin > 30) {
+      VLOG(1) << "Found tag number " << det->id << " hamming: " << det->hamming
+              << " margin: " << det->decision_margin;
+
+      const aos::monotonic_clock::time_point before_pose_estimation =
+          aos::monotonic_clock::now();
+      // First create an apriltag_detection_info_t struct using your known
+      // parameters.
+      apriltag_detection_info_t info;
+      info.det = det;
+      info.tagsize = 0.1524;
+      info.fx = intrinsics_.at<double>(0, 0);
+      info.fy = intrinsics_.at<double>(1, 1);
+      info.cx = intrinsics_.at<double>(0, 2);
+      info.cy = intrinsics_.at<double>(1, 2);
+
+      apriltag_pose_t pose;
+      double err = estimate_tag_pose(&info, &pose);
+
+      VLOG(1) << "err: " << err;
+
+      results.emplace_back(*det, pose);
+
+      const aos::monotonic_clock::time_point after_pose_estimation =
+          aos::monotonic_clock::now();
+
+      VLOG(1) << "Took "
+              << std::chrono::duration<double>(after_pose_estimation -
+                                               before_pose_estimation)
+                     .count()
+              << " seconds for pose estimation";
+    }
+  }
+
+  apriltag_detections_destroy(detections);
+
+  const aos::monotonic_clock::time_point end_time = aos::monotonic_clock::now();
+
+  timeprofile_display(tag_detector_->tp);
+
+  VLOG(1) << "Took "
+          << std::chrono::duration<double>(end_time - start_time).count()
+          << " seconds to detect overall";
+
+  return results;
+}
+
+void AprilViewerMain() {
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  aos::ShmEventLoop event_loop(&config.message());
+
+  AprilRoboticsDetector detector(&event_loop, "/camera");
+
+  detector.SetWorkerpoolAffinities();
+
+  event_loop.SetRuntimeAffinity(aos::MakeCpusetFromCpus({5}));
+
+  struct sched_param param;
+  param.sched_priority = 21;
+  PCHECK(sched_setscheduler(0, SCHED_FIFO, &param) == 0);
+
+  event_loop.Run();
+}
+
+}  // namespace vision
+}  // namespace y2023
+
+int main(int argc, char **argv) {
+  aos::InitGoogle(&argc, &argv);
+  y2023::vision::AprilViewerMain();
+
+  return 0;
+}
diff --git a/y2023/vision/aprilrobotics.h b/y2023/vision/aprilrobotics.h
new file mode 100644
index 0000000..f36f3eb
--- /dev/null
+++ b/y2023/vision/aprilrobotics.h
@@ -0,0 +1,102 @@
+
+#include <string>
+
+#include "aos/events/event_loop.h"
+#include "aos/events/shm_event_loop.h"
+#include "aos/init.h"
+#include "aos/network/team_number.h"
+#include "aos/realtime.h"
+#include "frc971/vision/calibration_generated.h"
+#include "frc971/vision/charuco_lib.h"
+#include "frc971/vision/target_map_generated.h"
+#include "frc971/vision/target_mapper.h"
+#include "frc971/vision/vision_generated.h"
+#include "opencv2/core/eigen.hpp"
+#include "opencv2/imgproc.hpp"
+#include "third_party/apriltag/apriltag.h"
+#include "third_party/apriltag/apriltag_pose.h"
+#include "third_party/apriltag/tag16h5.h"
+#include "y2023/vision/calibration_data.h"
+
+DECLARE_int32(team_number);
+
+namespace y2023 {
+namespace vision {
+
+class AprilRoboticsDetector {
+ public:
+  AprilRoboticsDetector(aos::EventLoop *event_loop,
+                        std::string_view channel_name);
+
+  ~AprilRoboticsDetector();
+
+  void SetWorkerpoolAffinities();
+
+  std::vector<std::pair<apriltag_detection_t, apriltag_pose_t>> DetectTags(
+      cv::Mat image);
+
+ private:
+  void HandleImage(cv::Mat image);
+
+  flatbuffers::Offset<frc971::vision::TargetPoseFbs> BuildTargetPose(
+      const apriltag_pose_t &pose,
+      frc971::vision::TargetMapper::TargetId target_id,
+      flatbuffers::FlatBufferBuilder *fbb);
+
+  static const frc971::vision::calibration::CameraCalibration *
+  FindCameraCalibration(
+      const frc971::vision::calibration::CalibrationData *calibration_data,
+      std::string_view node_name) {
+    for (const frc971::vision::calibration::CameraCalibration *candidate :
+         *calibration_data->camera_calibrations()) {
+      if (candidate->node_name()->string_view() != node_name) {
+        continue;
+      }
+      if (candidate->team_number() != FLAGS_team_number) {
+        continue;
+      }
+      return candidate;
+    }
+    LOG(FATAL) << ": Failed to find camera calibration for " << node_name
+               << " on " << FLAGS_team_number;
+  }
+
+  static cv::Mat CameraIntrinsics(
+      const frc971::vision::calibration::CameraCalibration
+          *camera_calibration) {
+    cv::Mat result(3, 3, CV_32F,
+                   const_cast<void *>(static_cast<const void *>(
+                       camera_calibration->intrinsics()->data())));
+    result.convertTo(result, CV_64F);
+    CHECK_EQ(result.total(), camera_calibration->intrinsics()->size());
+
+    return result;
+  }
+
+  static cv::Mat CameraDistCoeffs(
+      const frc971::vision::calibration::CameraCalibration
+          *camera_calibration) {
+    const cv::Mat result(5, 1, CV_32F,
+                         const_cast<void *>(static_cast<const void *>(
+                             camera_calibration->dist_coeffs()->data())));
+    CHECK_EQ(result.total(), camera_calibration->dist_coeffs()->size());
+    return result;
+  }
+
+  apriltag_family_t *tag_family_;
+  apriltag_detector_t *tag_detector_;
+
+  const aos::FlatbufferSpan<frc971::vision::calibration::CalibrationData>
+      calibration_data_;
+  const frc971::vision::calibration::CameraCalibration *calibration_;
+  cv::Mat intrinsics_;
+  cv::Mat camera_distortion_coeffs_;
+
+  aos::Ftrace ftrace_;
+
+  frc971::vision::ImageCallback image_callback_;
+  aos::Sender<frc971::vision::TargetMap> target_map_sender_;
+};
+
+}  // namespace vision
+}  // namespace y2023
diff --git a/y2023/y2023_pi_template.json b/y2023/y2023_pi_template.json
index fd3e97f..4c14031 100644
--- a/y2023/y2023_pi_template.json
+++ b/y2023/y2023_pi_template.json
@@ -167,7 +167,63 @@
       "max_size": 4200000,
       "num_readers": 4,
       "read_method": "PIN",
-      "num_senders": 18
+      "num_senders": 18,
+      "logger": "NOT_LOGGED"
+    },
+    {
+      "name": "/pi{{ NUM }}/camera/decimated",
+      "type": "frc971.vision.CameraImage",
+      "source_node": "pi{{ NUM }}",
+      "frequency": 2,
+      "max_size": 4200000,
+      "num_senders": 2
+    },
+    {
+      "name": "/pi{{ NUM }}/camera",
+      "type": "frc971.vision.TargetMap",
+      "source_node": "pi{{ NUM }}",
+      "frequency": 80,
+      "num_senders": 2,
+      "max_size": 40000,
+      "logger": "LOCAL_AND_REMOTE_LOGGER",
+      "logger_nodes": [
+        "imu",
+        "logger"
+      ],
+      "destination_nodes": [
+        {
+          "name": "imu",
+          "priority": 4,
+          "timestamp_logger": "LOCAL_AND_REMOTE_LOGGER",
+          "timestamp_logger_nodes": [
+            "pi{{ NUM }}"
+          ],
+          "time_to_live": 5000000
+        },
+        {
+          "name": "logger",
+          "priority": 4,
+          "timestamp_logger": "LOCAL_AND_REMOTE_LOGGER",
+          "timestamp_logger_nodes": [
+            "pi{{ NUM }}"
+          ],
+          "time_to_live": 5000000
+        }
+      ]
+    },
+    {
+      "name": "/pi{{ NUM }}/aos/remote_timestamps/imu/pi{{ NUM }}/camera/frc971-vision-TargetMap",
+      "type": "aos.message_bridge.RemoteMessage",
+      "frequency": 80,
+      "source_node": "pi{{ NUM }}",
+      "max_size": 208
+    },
+    {
+      "name": "/pi{{ NUM }}/aos/remote_timestamps/logger/pi{{ NUM }}/camera/frc971-vision-TargetMap",
+      "type": "aos.message_bridge.RemoteMessage",
+      "frequency": 80,
+      "source_node": "pi{{ NUM }}",
+      "max_size": 208
     },
     {
       "name": "/logger/aos",