Created vision node

toid_vision/src/toid_vision.cpp

@@ -0,0 +1,177 @@
+#include "toid_vision/toid_vision.hpp"
+
+#include <string>
+#include <vector>
+
+namespace toid
+{
+
+NutDetector::NutDetector(rclcpp::NodeOptions & options)
+: Node("compressed_image_subscriber", options)
+{
+  detectorParams_ = cv::aruco::DetectorParameters::create();
+  dictionary_ = cv::aruco::getPredefinedDictionary(cv::aruco::DICT_4X4_100);
+
+  bool use_compressed = true;
+  this->declare_parameter("use_compressed", rclcpp::ParameterValue(true));
+  this->get_parameter("use_compressed", use_compressed);
+
+  this->declare_parameter("camera_frame_id", rclcpp::ParameterValue("camera"));
+  this->get_parameter("camera_frame_id", camera_frame_id_);
+
+  this->declare_parameter("is_blue", rclcpp::ParameterValue(true));
+  this->get_parameter("is_blue", is_blue_);
+
+  compressed_img_sub_.subscribe(this, "/camera/image_raw/compressed");
+  cam_info_sub_.subscribe(this, "/camera/camera_info");
+
+  sync_ = std::make_shared<message_filters::Synchronizer<SyncPolicy>>(
+    SyncPolicy(2), compressed_img_sub_, cam_info_sub_);
+
+  using namespace std::placeholders;
+  sync_->registerCallback(std::bind(&NutDetector::compressed_topic_callback, this, _1, _2));
+  RCLCPP_INFO(this->get_logger(), "NutDetector activated... scanning for DEEZ NUTZ");
+
+  pose_pub_ =
+    this->create_publisher<geometry_msgs::msg::PoseStamped>("/closest_acorn", rclcpp::QoS(1));
+flip_pub_ = 
+    this->create_publisher<std_msgs::msg::String>("/to_flip", rclcpp::QoS(1));
+}
+
+void NutDetector::compressed_topic_callback(
+  const sensor_msgs::msg::CompressedImage::ConstSharedPtr & msg,
+  const sensor_msgs::msg::CameraInfo::ConstSharedPtr & cam_info)
+{
+  try {
+    cv::Mat rawData(1, msg->data.size(), CV_8UC1, (void *)msg->data.data());
+    cv::Mat decodedImage = cv::imdecode(rawData, cv::IMREAD_COLOR);
+
+    if (decodedImage.empty()) {
+      RCLCPP_ERROR(this->get_logger(), "Could not decode image!");
+      return;
+    }
+
+    process_image(decodedImage, msg->header, cam_info);
+
+  } catch (const std::exception & e) {
+    RCLCPP_ERROR(this->get_logger(), "Error processing image: %s", e.what());
+  }
+}
+
+void NutDetector::process_image(
+  cv::Mat & decodedImage, const std_msgs::msg::Header & header,
+  const sensor_msgs::msg::CameraInfo::ConstSharedPtr & cam_info)
+{
+  markers_.clear();
+
+  extract_markers(decodedImage);
+
+  if (markers_.empty()) {
+    return;
+  }
+
+  std::vector<cv::Vec3d> rvecs, tvecs;
+
+  cv::Mat k = cv::Mat(3, 3, CV_64F, (void *)cam_info->k.data());
+  cv::aruco::estimatePoseSingleMarkers(markerCorners_, 0.03, k, cam_info->d, rvecs, tvecs);
+
+  size_t prev = 0;
+  std::string colors = "";
+
+  for (size_t i = 0; i < markerIds_.size(); i++) {
+    const int idx = markers_[i].second;
+    if (markerIds_[idx] != 47 && markerIds_[idx] != 36) {
+      continue;
+    }
+
+    RCLCPP_DEBUG(
+      this->get_logger(), "Marker %ld at %lf %lf", i, markerCorners_[idx][0].x,
+      markerCorners_[idx][0].y);
+
+    RCLCPP_DEBUG(
+      this->get_logger(), "World coords: %lf %lf %lf", tvecs[idx][0], tvecs[idx][1], tvecs[idx][2]);
+
+    float dx = tvecs[idx][0] - tvecs[prev][0];
+    float dy = tvecs[idx][1] - tvecs[prev][1];
+    float dz = tvecs[idx][2] - tvecs[prev][2];
+
+    float dist = std::sqrt(dx * dx + dy * dy + dz * dz);
+
+    RCLCPP_DEBUG(this->get_logger(), "Dist from prev: %lf", dist);
+
+    if (dist > 0.07) {
+      continue;
+    }
+
+    prev = idx;
+
+    colors += (markerIds_[idx] == 36) ? "b" : "y";
+
+    if (colors.size() > 3) {
+      break;
+    }
+  }
+
+  if (colors.size() == 4) {
+    RCLCPP_DEBUG(this->get_logger(), "Last seen orientation: %s", colors.c_str());
+    for(int i = 0; i<4; i++) {
+        if(colors[i] == 'b') {
+            colors[i] = (is_blue_)? '0' : '1';
+        } else {
+            colors[i] = (is_blue_)? '1' : '0';
+        }
+    }
+
+    RCLCPP_DEBUG(this->get_logger(), "Last seen orientation: %s", colors.c_str());
+    std_msgs::msg::String m;
+    m.data = colors;
+    flip_pub_->publish(m);
+  }
+
+  geometry_msgs::msg::PoseStamped pose;
+  get_pose(tvecs[markers_[0].second], rvecs[markers_[0].second], pose.pose);
+  pose.header.frame_id = camera_frame_id_;
+  pose.header.stamp = header.stamp;
+  pose_pub_->publish(pose);
+}
+
+void NutDetector::extract_markers(cv::Mat & image)
+{
+  cv::aruco::detectMarkers(image, dictionary_, markerCorners_, markerIds_, detectorParams_);
+
+  // cv::aruco::drawDetectedMarkers(image, markerCorners_, markerIds_);
+
+  markers_.reserve(markerIds_.size());
+  markers_.resize(markerIds_.size());
+
+  for (size_t i = 0; i < markerIds_.size(); i++) {
+    markers_[i].first = markerCorners_[i];
+    markers_[i].second = i;
+  }
+
+  std::sort(
+    markers_.begin(), markers_.end(),
+    [](std::pair<std::vector<cv::Point2f>, int> & a, std::pair<std::vector<cv::Point2f>, int> & b) {
+      return a.first[0].y > b.first[0].y;
+    });
+}
+
+void NutDetector::get_pose(cv::Vec3d & tvec, cv::Vec3d & rvec, geometry_msgs::msg::Pose & out_pose)
+{
+  cv::Mat R;
+  cv::Rodrigues(rvec, R);
+  tf2::Matrix3x3 tf2_rot{R.at<double>(0, 0), R.at<double>(0, 1), R.at<double>(0, 2),
+                         R.at<double>(1, 0), R.at<double>(1, 1), R.at<double>(1, 2),
+                         R.at<double>(2, 0), R.at<double>(2, 1), R.at<double>(2, 2)};
+
+  tf2::Quaternion q;
+  tf2_rot.getRotation(q);
+
+  tf2::convert(q, out_pose.orientation);
+
+  out_pose.position.x = tvec[0];
+  out_pose.position.y = tvec[1];
+  out_pose.position.z = tvec[2];
+}
+
+}  // namespace toid

toid_vision/src/vision.cpp

@@ -0,0 +1,10 @@
+#include "toid_vision/toid_vision.hpp"
+
+int main(int argc, char * argv[])
+{
+  rclcpp::init(argc, argv);
+  rclcpp::NodeOptions options;
+  rclcpp::spin(std::make_shared<toid::NutDetector>(options));
+  rclcpp::shutdown();
+  return 0;
+}