Store target map in a flatbuffer

Mapping process takes input json with initial guesses for target poses,
and outputs a json with the solved poses.

Signed-off-by: Milind Upadhyay <milind.upadhyay@gmail.com>
Change-Id: I29b82aacd37758e7dc8d8cd383821899182a1950
diff --git a/frc971/vision/BUILD b/frc971/vision/BUILD
index 7872926..54dab0d 100644
--- a/frc971/vision/BUILD
+++ b/frc971/vision/BUILD
@@ -88,6 +88,14 @@
     ],
 )
 
+flatbuffer_cc_library(
+    name = "target_map_fbs",
+    srcs = ["target_map.fbs"],
+    gen_reflections = 1,
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "target_mapper",
     srcs = ["target_mapper.cc"],
@@ -96,6 +104,7 @@
     visibility = ["//visibility:public"],
     deps = [
         ":geometry_lib",
+        ":target_map_fbs",
         "//aos/events:simulated_event_loop",
         "//frc971/control_loops:control_loop",
         "//frc971/vision/ceres:pose_graph_2d_lib",
diff --git a/frc971/vision/target_map.fbs b/frc971/vision/target_map.fbs
new file mode 100644
index 0000000..8adaaaa
--- /dev/null
+++ b/frc971/vision/target_map.fbs
@@ -0,0 +1,24 @@
+namespace frc971.vision;
+
+// Represents 3d pose of an april tag on the field.
+table TargetPoseFbs {
+  // AprilTag ID of this target
+  id:uint64 (id: 0);
+
+  // Pose of target relative to field origin.
+  // NOTE: As of now, we only solve for the 2d pose (x, y, yaw)
+  // and all other values will be 0.
+  x:double (id: 1);
+  y:double (id: 2);
+  z:double (id: 3);
+
+  roll:double (id: 4);
+  pitch:double (id: 5);
+  yaw:double (id: 6);
+}
+
+// Map of all target poses on a field.
+// This would be solved for by TargetMapper
+table TargetMap {
+  target_poses:[TargetPoseFbs] (id: 0);
+}
\ No newline at end of file
diff --git a/frc971/vision/target_mapper.cc b/frc971/vision/target_mapper.cc
index b806c03..3342670 100644
--- a/frc971/vision/target_mapper.cc
+++ b/frc971/vision/target_mapper.cc
@@ -207,6 +207,18 @@
 }
 
 TargetMapper::TargetMapper(
+    std::string_view target_poses_path,
+    std::vector<ceres::examples::Constraint2d> target_constraints)
+    : target_constraints_(target_constraints) {
+  aos::FlatbufferDetachedBuffer<TargetMap> target_map =
+      aos::JsonFileToFlatbuffer<TargetMap>(target_poses_path);
+  for (const auto *target_pose_fbs : *target_map.message().target_poses()) {
+    target_poses_[target_pose_fbs->id()] = ceres::examples::Pose2d{
+        target_pose_fbs->x(), target_pose_fbs->y(), target_pose_fbs->yaw()};
+  }
+}
+
+TargetMapper::TargetMapper(
     std::map<TargetId, ceres::examples::Pose2d> target_poses,
     std::vector<ceres::examples::Constraint2d> target_constraints)
     : target_poses_(target_poses), target_constraints_(target_constraints) {}
@@ -223,15 +235,6 @@
 }
 
 // Taken from ceres/examples/slam/pose_graph_2d/pose_graph_2d.cc
-void TargetMapper::OutputPoses(
-    const std::map<int, ceres::examples::Pose2d> &poses) {
-  for (const auto &pair : poses) {
-    std::cout << pair.first << ": " << pair.second.x << " " << pair.second.y
-              << ' ' << pair.second.yaw_radians << std::endl;
-  }
-}
-
-// Taken from ceres/examples/slam/pose_graph_2d/pose_graph_2d.cc
 void TargetMapper::BuildOptimizationProblem(
     std::map<int, ceres::examples::Pose2d> *poses,
     const std::vector<ceres::examples::Constraint2d> &constraints,
@@ -306,19 +309,49 @@
   ceres::Solver::Summary summary;
   ceres::Solve(options, problem, &summary);
 
-  std::cout << summary.FullReport() << '\n';
+  LOG(INFO) << summary.FullReport() << '\n';
 
   return summary.IsSolutionUsable();
 }
 
-void TargetMapper::Solve() {
+void TargetMapper::Solve(std::optional<std::string_view> output_path) {
   ceres::Problem problem;
   BuildOptimizationProblem(&target_poses_, target_constraints_, &problem);
 
   CHECK(SolveOptimizationProblem(&problem))
       << "The solve was not successful, exiting.";
 
-  OutputPoses(target_poses_);
+  // TODO(milind): add origin to first target offset to all poses
+
+  auto map_json = MapToJson();
+  VLOG(1) << "Solved target poses: " << map_json;
+  if (output_path.has_value()) {
+    LOG(INFO) << "Writing map to file: " << output_path.value();
+    aos::util::WriteStringToFileOrDie(output_path.value(), map_json);
+  }
+}
+
+std::string TargetMapper::MapToJson() const {
+  flatbuffers::FlatBufferBuilder fbb;
+
+  // Convert poses to flatbuffers
+  std::vector<flatbuffers::Offset<TargetPoseFbs>> target_poses_fbs;
+  for (const auto &[id, pose] : target_poses_) {
+    TargetPoseFbs::Builder target_pose_builder(fbb);
+    target_pose_builder.add_id(id);
+    target_pose_builder.add_x(id);
+    target_pose_builder.add_y(id);
+    target_pose_builder.add_yaw(id);
+
+    target_poses_fbs.emplace_back(target_pose_builder.Finish());
+  }
+
+  flatbuffers::Offset<TargetMap> target_map_offset =
+      CreateTargetMap(fbb, fbb.CreateVector(target_poses_fbs));
+
+  return aos::FlatbufferToJson(
+      flatbuffers::GetMutableTemporaryPointer(fbb, target_map_offset),
+      {.multi_line = true});
 }
 
 }  // namespace frc971::vision
diff --git a/frc971/vision/target_mapper.h b/frc971/vision/target_mapper.h
index 55d3c7e..df9f5b4 100644
--- a/frc971/vision/target_mapper.h
+++ b/frc971/vision/target_mapper.h
@@ -5,6 +5,7 @@
 
 #include "aos/events/simulated_event_loop.h"
 #include "frc971/vision/ceres/types.h"
+#include "frc971/vision/target_map_generated.h"
 
 namespace frc971::vision {
 
@@ -18,17 +19,27 @@
 
   struct TargetPose {
     TargetId id;
+    // TOOD(milind): switch everything to 3d once we're more confident in 2d
+    // solving
     ceres::examples::Pose2d pose;
   };
 
-  // target_poses are initial guesses for the actual locations of the targets on
-  // the field.
+  // target_poses_path is the path to a TargetMap json with initial guesses for
+  // the actual locations of the targets on the field.
   // target_constraints are the deltas between consecutive target detections,
   // and are usually prepared by the DataAdapter class below.
+  TargetMapper(std::string_view target_poses_path,
+               std::vector<ceres::examples::Constraint2d> target_constraints);
+  // Alternate constructor for tests.
+  // Takes in the actual intial guesses instead of a file containing them
   TargetMapper(std::map<TargetId, ceres::examples::Pose2d> target_poses,
                std::vector<ceres::examples::Constraint2d> target_constraints);
 
-  void Solve();
+  // If output_path is set, the map will be saved to that file as a json
+  void Solve(std::optional<std::string_view> output_path = std::nullopt);
+
+  // Prints target poses into a TargetMap flatbuffer json
+  std::string MapToJson() const;
 
   static std::optional<TargetPose> GetTargetPoseById(
       std::vector<TargetPose> target_poses, TargetId target_id);
@@ -38,9 +49,6 @@
   }
 
  private:
-  // Output the poses to std::cout with format: ID: x y yaw_radians.
-  static void OutputPoses(const std::map<int, ceres::examples::Pose2d> &poses);
-
   // Constructs the nonlinear least squares optimization problem from the
   // pose graph constraints.
   void BuildOptimizationProblem(