Added H26x decoding node for femto_mega

2026-10-07 05:27:45 +08:00 · 2024-09-18 22:21:04 +08:00
parent cdd1c9fbd9
commit 3cf375fb07
6 changed files with 342 additions and 113 deletions
@@ -0,0 +1,145 @@
+extern "C" {
+#include <libavcodec/avcodec.h>
+#include <libavformat/avformat.h>
+#include <libswscale/swscale.h>
+#include <libavutil/imgutils.h>
+#include <libavutil/time.h>
+}
+#include <memory>
+
+#include "rclcpp/rclcpp.hpp"
+#include "orbbec_camera/ob_camera_node.h"
+#include <thread>
+#include <geometry_msgs/msg/transform_stamped.hpp>
+
+#include "orbbec_camera/utils.h"
+#include <filesystem>
+#include <fstream>
+#include "diagnostic_msgs/msg/diagnostic_status.hpp"
+#include "libobsensor/hpp/Utils.hpp"
+
+class H264DecoderNode : public rclcpp::Node {
+ public:
+  H264DecoderNode() : Node("h264_decoder_node") {
+    avformat_network_init();
+    rclcpp::QoS qos_settings(30);
+    qos_settings.reliability(RMW_QOS_POLICY_RELIABILITY_BEST_EFFORT);
+    qos_settings.durability(RMW_QOS_POLICY_DURABILITY_VOLATILE);
+    qos_settings.history(RMW_QOS_POLICY_HISTORY_KEEP_LAST);
+
+    compressed_image_subscriber_ = this->create_subscription<sensor_msgs::msg::CompressedImage>(
+        "/camera/color/h26x_encoded_data", qos_settings,
+        std::bind(&H264DecoderNode::compressedImageCallback, this, std::placeholders::_1));
+
+    rgb_image_publisher_ = this->create_publisher<sensor_msgs::msg::Image>(
+        "/camera/color/h26x_decoder/image_raw", qos_settings);
+  }
+
+ private:
+  void decode_init(const sensor_msgs::msg::CompressedImage::SharedPtr msg) {
+    if (msg->format == "h264") {
+      codec_ = std::shared_ptr<const AVCodec>(avcodec_find_decoder(AV_CODEC_ID_H264),
+                                              [](const AVCodec*) {});
+      codec_context_ = std::shared_ptr<AVCodecContext>(avcodec_alloc_context3(codec_.get()),
+                                                       [](AVCodecContext* ctx) {
+                                                         if (ctx) {
+                                                           avcodec_free_context(&ctx);
+                                                         }
+                                                       });
+
+      if (avcodec_open2(codec_context_.get(), codec_.get(), nullptr) < 0) {
+        RCLCPP_ERROR(this->get_logger(), "Failed to open codec");
+        return;
+      }
+
+      frame_ = std::shared_ptr<AVFrame>(av_frame_alloc(), [](AVFrame* f) {
+        if (f) av_frame_free(&f);
+      });
+
+      packet_ = std::shared_ptr<AVPacket>(av_packet_alloc(), [](AVPacket* p) {
+        if (p) av_packet_free(&p);
+      });
+      codec_init_ = 0;
+    } else if (msg->format == "h265") {
+      codec_ = std::shared_ptr<const AVCodec>(avcodec_find_decoder(AV_CODEC_ID_HEVC),
+                                              [](const AVCodec*) {});
+      codec_context_ = std::shared_ptr<AVCodecContext>(avcodec_alloc_context3(codec_.get()),
+                                                       [](AVCodecContext* ctx) {
+                                                         if (ctx) {
+                                                           avcodec_free_context(&ctx);
+                                                         }
+                                                       });
+
+      if (avcodec_open2(codec_context_.get(), codec_.get(), nullptr) < 0) {
+        RCLCPP_ERROR(this->get_logger(), "Failed to open codec");
+        return;
+      }
+
+      frame_ = std::shared_ptr<AVFrame>(av_frame_alloc(), [](AVFrame* f) {
+        if (f) av_frame_free(&f);
+      });
+
+      packet_ = std::shared_ptr<AVPacket>(av_packet_alloc(), [](AVPacket* p) {
+        if (p) av_packet_free(&p);
+      });
+      codec_init_ = 0;
+    }
+  }
+  void decode_frame(const sensor_msgs::msg::CompressedImage::SharedPtr msg) {
+    av_packet_unref(packet_.get());
+    packet_->data = const_cast<uint8_t*>(msg->data.data());
+    packet_->size = msg->data.size();
+    std::stringstream ss;
+    const size_t bytes_to_print = std::min<size_t>(msg->data.size(), 32);
+    for (size_t i = 0; i < bytes_to_print; ++i) {
+      ss << std::hex << std::setw(2) << std::setfill('0') << static_cast<int>(msg->data[i]) << " ";
+    }
+    RCLCPP_INFO(this->get_logger(), "Data (hex): %s", ss.str().c_str());
+
+    send_ret_ = avcodec_send_packet(codec_context_.get(), packet_.get());
+    if (send_ret_ >= 0) {
+      receive_ret_ = avcodec_receive_frame(codec_context_.get(), frame_.get());
+      if (receive_ret_ >= 0) {
+        cv::Mat rgb_image(frame_->height, frame_->width, CV_8UC3);
+        SwsContext* sws_context = sws_getContext(frame_->width, frame_->height,
+                                                 static_cast<AVPixelFormat>(frame_->format),
+                                                 frame_->width, frame_->height, AV_PIX_FMT_BGR24,
+                                                 SWS_BILINEAR, nullptr, nullptr, nullptr);
+
+        uint8_t* dest[1] = {rgb_image.data};
+        int linesize[1] = {static_cast<int>(rgb_image.step1())};
+        sws_scale(sws_context, frame_->data, frame_->linesize, 0, frame_->height, dest, linesize);
+        sws_freeContext(sws_context);
+        auto rgb_image_msg =
+            cv_bridge::CvImage(std_msgs::msg::Header(), "bgr8", rgb_image).toImageMsg();
+        rgb_image_msg->header.stamp = this->now();
+
+        rgb_image_publisher_->publish(*rgb_image_msg);
+      }
+    }
+  }
+  void compressedImageCallback(const sensor_msgs::msg::CompressedImage::SharedPtr msg) {
+    RCLCPP_INFO(this->get_logger(), "Format: %s", msg->format.c_str());
+    if (codec_init_) {
+      decode_init(msg);
+    }
+    decode_frame(msg);
+  }
+
+  rclcpp::Subscription<sensor_msgs::msg::CompressedImage>::SharedPtr compressed_image_subscriber_;
+  rclcpp::Publisher<sensor_msgs::msg::Image>::SharedPtr rgb_image_publisher_;
+  std::shared_ptr<const AVCodec> codec_;
+  std::shared_ptr<AVCodecContext> codec_context_;
+  std::shared_ptr<AVFrame> frame_;
+  std::shared_ptr<AVPacket> packet_;
+  int codec_init_ = 1;
+  int send_ret_;
+  int receive_ret_;
+};
+int main(int argc, char** argv) {
+  rclcpp::init(argc, argv);
+  auto node = std::make_shared<H264DecoderNode>();
+  rclcpp::spin(node);
+  rclcpp::shutdown();
+  return 0;
+}