Skip to content

Commit

Permalink
feat(tier4_screen_capture_rviz_plugin): add video capture with buffer…
Browse files Browse the repository at this point in the history 
… 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 <[email protected]>


---------

Signed-off-by: kyoichi-sugahara <[email protected]>
  • Loading branch information
@kyoichi-sugahara
kyoichi-sugahara authored Jul 25, 2024
1 parent 3baf48f commit 97ed6ee
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 16 deletions.
9 changes: 5 additions & 4 deletions common/tier4_screen_capture_rviz_plugin/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
89 changes: 78 additions & 11 deletions common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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_);
}
Expand All @@ -105,6 +114,9 @@ void AutowareScreenCapturePanel::onInitialize()
capture_video_srv_ = raw_node_->create_service<std_srvs::srv::Trigger>(
"/debug/capture/video",
std::bind(&AutowareScreenCapturePanel::onCaptureVideoTrigger, this, _1, _2));
capture_video_with_buffer_srv_ = raw_node_->create_service<std_srvs::srv::Trigger>(
"/debug/capture/video_with_buffer",
std::bind(&AutowareScreenCapturePanel::onCaptureVideoWithBufferTrigger, this, _1, _2));
capture_screen_shot_srv_ = raw_node_->create_service<std_srvs::srv::Trigger>(
"/debug/capture/screen_shot",
std::bind(&AutowareScreenCapturePanel::onCaptureScreenShotTrigger, this, _1, _2));
Expand All @@ -126,7 +138,7 @@ void AutowareScreenCapturePanel::onClickScreenCapture()
time_text + ".png");
}

void AutowareScreenCapturePanel::onClickVideoCapture()
void AutowareScreenCapturePanel::onClickVideoCapture(bool use_buffer)
{
const int clock = static_cast<int>(1e3 / capture_hz_->value());
try {
Expand Down Expand Up @@ -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;");
Expand All @@ -189,16 +210,62 @@ void AutowareScreenCapturePanel::onTimer()
cv::Mat image(
size, CV_8UC3, const_cast<uchar *>(q_image.bits()),
static_cast<size_t>(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<uchar *>(q_image.bits()),
static_cast<size_t>(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());
Expand Down
15 changes: 14 additions & 1 deletion common/tier4_screen_capture_rviz_plugin/src/screen_capture_panel.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,10 @@
// ros
#include <std_srvs/srv/trigger.hpp>

#include <deque>
#include <memory>
#include <string>
#include <utility>
#include <vector>

class QLineEdit;
Expand All @@ -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();

Expand All @@ -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<cv::Mat> image_vec_;
std::deque<std::pair<cv::Mat, rclcpp::Time>> 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)
{
Expand All @@ -102,6 +114,7 @@ public Q_SLOTS:

protected:
rclcpp::Service<std_srvs::srv::Trigger>::SharedPtr capture_video_srv_;
rclcpp::Service<std_srvs::srv::Trigger>::SharedPtr capture_video_with_buffer_srv_;
rclcpp::Service<std_srvs::srv::Trigger>::SharedPtr capture_screen_shot_srv_;
rclcpp::Node::SharedPtr raw_node_;
};
Expand Down

0 comments on commit 97ed6ee

Please sign in to comment.