From 97ed6ee619625a60895b05a3a22042da243af780 Mon Sep 17 00:00:00 2001 From: Kyoichi Sugahara Date: Thu, 25 Jul 2024 17:50:38 +0900 Subject: [PATCH] feat(tier4_screen_capture_rviz_plugin): add video capture with buffer functionality (#80) * feat: add video capture with buffer functionality The code changes introduce new functions for capturing video with a buffer. This allows for capturing and saving video frames with a delay, enabling the user to save a specified number of frames in memory before writing them to a file. Signed-off-by: kyoichi-sugahara --------- Signed-off-by: kyoichi-sugahara --- .../README.md | 9 +- .../src/screen_capture_panel.cpp | 89 ++++++++++++++++--- .../src/screen_capture_panel.hpp | 15 +++- 3 files changed, 97 insertions(+), 16 deletions(-) diff --git a/common/tier4_screen_capture_rviz_plugin/README.md b/common/tier4_screen_capture_rviz_plugin/README.md index d16c19c3..40591c9d 100644 --- a/common/tier4_screen_capture_rviz_plugin/README.md +++ b/common/tier4_screen_capture_rviz_plugin/README.md @@ -6,10 +6,11 @@ This plugin captures the screen of rviz. ## Interface -| Name | Type | Description | -| ---------------------------- | ------------------------ | ---------------------------------- | -| `/debug/capture/video` | `std_srvs::srv::Trigger` | Trigger to start screen capturing. | -| `/debug/capture/screen_shot` | `std_srvs::srv::Trigger` | Trigger to capture screen shot. | +| Name | Type | Description | +| ---------------------------------- | ------------------------ | ---------------------------------------------- | +| `/debug/capture/video` | `std_srvs::srv::Trigger` | Trigger to start screen capturing. | +| `/debug/capture/video_with_buffer` | `std_srvs::srv::Trigger` | Trigger to start screen capturing with buffer. | +| `/debug/capture/screen_shot` | `std_srvs::srv::Trigger` | Trigger to capture screen shot. | ## Assumptions / Known limits diff --git a/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.cpp b/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.cpp index cad82890..a64d71a3 100644 --- a/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.cpp +++ b/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.cpp @@ -85,7 +85,16 @@ void AutowareScreenCapturePanel::onCaptureVideoTrigger( [[maybe_unused]] const std_srvs::srv::Trigger::Request::SharedPtr req, const std_srvs::srv::Trigger::Response::SharedPtr res) { - onClickVideoCapture(); + onClickVideoCapture(false); + res->success = true; + res->message = stateToString(state_); +} + +void AutowareScreenCapturePanel::onCaptureVideoWithBufferTrigger( + [[maybe_unused]] const std_srvs::srv::Trigger::Request::SharedPtr req, + const std_srvs::srv::Trigger::Response::SharedPtr res) +{ + onClickVideoCapture(true); res->success = true; res->message = stateToString(state_); } @@ -105,6 +114,9 @@ void AutowareScreenCapturePanel::onInitialize() capture_video_srv_ = raw_node_->create_service( "/debug/capture/video", std::bind(&AutowareScreenCapturePanel::onCaptureVideoTrigger, this, _1, _2)); + capture_video_with_buffer_srv_ = raw_node_->create_service( + "/debug/capture/video_with_buffer", + std::bind(&AutowareScreenCapturePanel::onCaptureVideoWithBufferTrigger, this, _1, _2)); capture_screen_shot_srv_ = raw_node_->create_service( "/debug/capture/screen_shot", std::bind(&AutowareScreenCapturePanel::onCaptureScreenShotTrigger, this, _1, _2)); @@ -126,7 +138,7 @@ void AutowareScreenCapturePanel::onClickScreenCapture() time_text + ".png"); } -void AutowareScreenCapturePanel::onClickVideoCapture() +void AutowareScreenCapturePanel::onClickVideoCapture(bool use_buffer) { const int clock = static_cast(1e3 / capture_hz_->value()); try { @@ -158,15 +170,24 @@ void AutowareScreenCapturePanel::onClickVideoCapture() .rgbSwapped() .size(); current_movie_size_ = cv::Size(q_size.width(), q_size.height()); - writer_.open( - "capture/" + file_name_prefix_->text().toStdString() + capture_file_name_ + ".mp4", - fourcc, capture_hz_->value(), current_movie_size_); + is_buffering_ = use_buffer; + if (!is_buffering_) { + writer_.open( + "capture/" + file_name_prefix_->text().toStdString() + capture_file_name_ + ".mp4", + fourcc, capture_hz_->value(), current_movie_size_); + } else { + frame_buffer_.clear(); + } } capture_timer_->start(clock); state_ = State::CAPTURING; break; case State::CAPTURING: - writer_.release(); + if (is_buffering_) { + saveBufferedVideo(); + } else { + writer_.release(); + } capture_timer_->stop(); capture_to_mp4_button_ptr_->setText("waiting for capture"); capture_to_mp4_button_ptr_->setStyleSheet("background-color: #00FF00;"); @@ -189,16 +210,62 @@ void AutowareScreenCapturePanel::onTimer() cv::Mat image( size, CV_8UC3, const_cast(q_image.bits()), static_cast(q_image.bytesPerLine())); - if (size != current_movie_size_) { - cv::Mat new_image; - cv::resize(image, new_image, current_movie_size_); - writer_.write(new_image); + + if (is_buffering_) { + // Buffering mode: Store frame in buffer + frame_buffer_.push_back({image.clone(), rclcpp::Clock().now()}); + if (frame_buffer_.size() > buffer_size_) { + frame_buffer_.pop_front(); + } } else { - writer_.write(image); + // Normal mode: Write frame to video file + if (size != current_movie_size_) { + cv::Mat new_image; + cv::resize(image, new_image, current_movie_size_); + writer_.write(new_image); + } else { + writer_.write(image); + } } + cv::waitKey(0); } +void AutowareScreenCapturePanel::captureFrameToBuffer() +{ + QScreen * screen = QGuiApplication::primaryScreen(); + QPixmap original_pixmap = screen->grabWindow(main_window_->winId()); + const auto q_image = + original_pixmap.toImage().convertToFormat(QImage::Format_RGB888).rgbSwapped(); + cv::Mat image( + q_image.height(), q_image.width(), CV_8UC3, const_cast(q_image.bits()), + static_cast(q_image.bytesPerLine())); + + frame_buffer_.push_back({image, rclcpp::Clock().now()}); + if (frame_buffer_.size() > buffer_size_) { + frame_buffer_.pop_front(); + } +} + +void AutowareScreenCapturePanel::saveBufferedVideo() +{ + if (frame_buffer_.empty()) return; + + int fourcc = cv::VideoWriter::fourcc('h', '2', '6', '4'); // mp4 + cv::VideoWriter buffered_writer; + buffered_writer.open( + "capture/" + file_name_prefix_->text().toStdString() + capture_file_name_ + "_buffered.mp4", + fourcc, capture_hz_->value(), current_movie_size_); + + for (const auto & [frame, _] : frame_buffer_) { + cv::Mat resized_frame; + cv::resize(frame, resized_frame, current_movie_size_); + buffered_writer.write(resized_frame); + } + + buffered_writer.release(); +} + void AutowareScreenCapturePanel::update() { setFormatDate(ros_time_label_, rclcpp::Clock().now().seconds()); diff --git a/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.hpp b/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.hpp index 5c4d16f5..78b90d4c 100644 --- a/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.hpp +++ b/common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.hpp @@ -41,8 +41,10 @@ // ros #include +#include #include #include +#include #include class QLineEdit; @@ -58,18 +60,23 @@ class AutowareScreenCapturePanel : public rviz_common::Panel void onInitialize() override; void createWallTimer(); void onTimer(); + void captureFrameToBuffer(); + void saveBufferedVideo(); void save(rviz_common::Config config) const override; void load(const rviz_common::Config & config) override; void onCaptureVideoTrigger( const std_srvs::srv::Trigger::Request::SharedPtr req, const std_srvs::srv::Trigger::Response::SharedPtr res); + void onCaptureVideoWithBufferTrigger( + const std_srvs::srv::Trigger::Request::SharedPtr req, + const std_srvs::srv::Trigger::Response::SharedPtr res); void onCaptureScreenShotTrigger( const std_srvs::srv::Trigger::Request::SharedPtr req, const std_srvs::srv::Trigger::Response::SharedPtr res); public Q_SLOTS: void onClickScreenCapture(); - void onClickVideoCapture(); + void onClickVideoCapture(bool use_buffer = false); void onPrefixChanged(); void onRateChanged(); @@ -85,9 +92,14 @@ public Q_SLOTS: State state_; std::string capture_file_name_; bool is_capture_{false}; + bool is_buffering_{false}; cv::VideoWriter writer_; cv::Size current_movie_size_; std::vector image_vec_; + std::deque> frame_buffer_; + // Size of the frame buffer (number of frames to keep in memory) + // At 10 Hz capture rate, 100 frames correspond to approximately 10 seconds of video + const size_t buffer_size_ = 100; static std::string stateToString(const State & state) { @@ -102,6 +114,7 @@ public Q_SLOTS: protected: rclcpp::Service::SharedPtr capture_video_srv_; + rclcpp::Service::SharedPtr capture_video_with_buffer_srv_; rclcpp::Service::SharedPtr capture_screen_shot_srv_; rclcpp::Node::SharedPtr raw_node_; };