y2022/vision/viewer_replay.cc - RealtimeRoboticsGroup/test - Gitiles

 #include "opencv2/calib3d.hpp"
 #include "opencv2/features2d.hpp"
 #include "opencv2/highgui/highgui.hpp"
 #include "opencv2/imgproc.hpp"

 #include "aos/events/logging/log_reader.h"
 #include "aos/events/simulated_event_loop.h"
 #include "aos/init.h"
 #include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
 #include "frc971/input/joystick_state_generated.h"
 #include "frc971/vision/vision_generated.h"
 #include "y2022/control_loops/superstructure/superstructure_status_generated.h"
 #include "y2022/vision/blob_detector.h"

 ABSL_FLAG(std::string, pi, "pi3", "Node name to replay.");
 ABSL_FLAG(std::string, image_save_prefix, "/tmp/img",
           "Prefix to use for saving images from the logfile.");
 ABSL_FLAG(bool, display, false, "If true, display the images with a timeout.");
 ABSL_FLAG(bool, detected_only, false,
           "If true, only write images which had blobs (unfiltered) detected");
 ABSL_FLAG(bool, filtered_only, false,
           "If true, only write images which had blobs (filtered) detected");
 ABSL_FLAG(bool, match_timestamps, false,
           "If true, name the files based on the time since the robot was "
           "enabled (match start). Only consider images during this time");
 ABSL_FLAG(std::string, logger_pi_log, "/tmp/logger_pi/",
           "Path to logger pi log");
 ABSL_FLAG(std::string, roborio_log, "/tmp/roborio/", "Path to roborio log");

 namespace y2022::vision {
 namespace {

 using aos::monotonic_clock;
 namespace superstructure = control_loops::superstructure;

 // Information to extract from the roborio log
 struct ReplayData {
   monotonic_clock::time_point match_start;
   monotonic_clock::time_point match_end;
   std::map<monotonic_clock::time_point, bool> robot_moving;
   std::map<monotonic_clock::time_point, superstructure::SuperstructureState>
       superstructure_states;
 };

 // Extract the useful data from the roborio log to be used for naming images
 void ReplayRoborio(ReplayData *data) {
   data->match_start = monotonic_clock::min_time;
   data->match_end = monotonic_clock::min_time;

   // Open logfiles
   aos::logger::LogReader reader(aos::logger::SortParts(
       aos::logger::FindLogs(absl::GetFlag(FLAGS_roborio_log))));
   reader.Register();
   const aos::Node *roborio =
       aos::configuration::GetNode(reader.configuration(), "roborio");

   std::unique_ptr<aos::EventLoop> event_loop =
       reader.event_loop_factory()->MakeEventLoop("roborio", roborio);

   auto joystick_state_fetcher =
       event_loop->MakeFetcher<aos::JoystickState>("/roborio/aos");
   auto drivetrain_status_fetcher =
       event_loop->MakeFetcher<frc971::control_loops::drivetrain::Status>(
           "/drivetrain");
   auto superstructure_status_fetcher =
       event_loop->MakeFetcher<superstructure::Status>("/superstructure");

   // Periodically check if the robot state updated
   event_loop->AddPhasedLoop(
       [&](int) {
         // Find the match start and end times
         if (joystick_state_fetcher.Fetch()) {
           if (data->match_start == monotonic_clock::min_time &&
               joystick_state_fetcher->enabled()) {
             data->match_start =
                 joystick_state_fetcher.context().monotonic_event_time;
           } else {
             if (data->match_end == monotonic_clock::min_time &&
                 data->match_start != monotonic_clock::min_time &&
                 !joystick_state_fetcher->autonomous() &&
                 !joystick_state_fetcher->enabled()) {
               data->match_end =
                   joystick_state_fetcher.context().monotonic_event_time;
             }
           }
         }

         // Add whether the robot was moving at a non-negligible speed to
         // the image name for debugging.
         drivetrain_status_fetcher.Fetch();
         if (drivetrain_status_fetcher.get()) {
           // If the robot speed was atleast this (m/s), it is
           // considered moving.
           constexpr double kMinMovingRobotSpeed = 0.5;
           data->robot_moving[drivetrain_status_fetcher.context()
                                  .monotonic_event_time] =
               (std::abs(drivetrain_status_fetcher->robot_speed()) >=
                kMinMovingRobotSpeed);
         }

         superstructure_status_fetcher.Fetch();
         if (superstructure_status_fetcher.get()) {
           data->superstructure_states[superstructure_status_fetcher.context()
                                           .monotonic_event_time] =
               superstructure_status_fetcher->state();
         }
       },
       std::chrono::milliseconds(50));
   reader.event_loop_factory()->Run();
 }

 template <typename T>
 T ClosestElement(const std::map<monotonic_clock::time_point, T> &map,
                  monotonic_clock::time_point now) {
   T closest;

   // The closest element is either the one right above it, or the element before
   // that one
   auto closest_it = map.lower_bound(now);
   // Handle the case where now is greater than all times in the map
   if (closest_it == map.cend()) {
     closest_it--;
     closest = closest_it->second;
   } else {
     // Start off with the closest as the first after now
     closest = closest_it->second;
     const monotonic_clock::duration after_duration = closest_it->first - now;
     closest_it--;

     // If there is a time before, check if that's closer to now
     if (closest_it != map.cbegin()) {
       const monotonic_clock::duration before_duration = now - closest_it->first;
       if (before_duration < after_duration) {
         closest = closest_it->second;
       }
     }
   }

   return closest;
 }

 // Extract images from the pi logs
 void ReplayPi(const ReplayData &data) {
   if (absl::GetFlag(FLAGS_match_timestamps)) {
     CHECK_NE(data.match_start, monotonic_clock::min_time)
         << "Can't use match timestamps if match never started";
     CHECK_NE(data.match_end, monotonic_clock::min_time)
         << "Can't use match timestamps if match never ended";
   }

   // Open logfiles
   aos::logger::LogReader reader(aos::logger::SortParts(
       aos::logger::FindLogs(absl::GetFlag(FLAGS_logger_pi_log))));
   reader.Register();
   const aos::Node *pi = aos::configuration::GetNode(reader.configuration(),
                                                     absl::GetFlag(FLAGS_pi));

   std::unique_ptr<aos::EventLoop> event_loop =
       reader.event_loop_factory()->MakeEventLoop("player", pi);

   LOG(INFO) << "Match start: " << data.match_start
             << ", match end: " << data.match_end;

   size_t nonmatch_image_count = 0;

   event_loop->MakeWatcher(
       "/camera/decimated", [&](const frc971::vision::CameraImage &image) {
         const auto match_start = data.match_start;
         // Find the closest robot moving and superstructure state to now
         const bool robot_moving =
             ClosestElement(data.robot_moving, event_loop->monotonic_now());
         const auto superstructure_state = ClosestElement(
             data.superstructure_states, event_loop->monotonic_now());

         if (absl::GetFlag(FLAGS_match_timestamps)) {
           if (event_loop->monotonic_now() < data.match_start) {
             // Ignore prematch images if we only care about ones during the
             // match
             return;
           } else if (event_loop->monotonic_now() >= data.match_end) {
             // We're done if the match is over and we only wanted match images
             reader.event_loop_factory()->Exit();
             return;
           }
         }

         // Create color image:
         cv::Mat image_color_mat(cv::Size(image.cols(), image.rows()), CV_8UC2,
                                 (void *)image.data()->data());
         cv::Mat image_mat(cv::Size(image.cols(), image.rows()), CV_8UC3);
         cv::cvtColor(image_color_mat, image_mat, cv::COLOR_YUV2BGR_YUYV);

         bool use_image = true;
         if (absl::GetFlag(FLAGS_detected_only) ||
             absl::GetFlag(FLAGS_filtered_only)) {
           // TODO(milind): if adding target estimation here in the future,
           // undistortion is needed
           BlobDetector::BlobResult blob_result;
           BlobDetector::ExtractBlobs(image_mat, &blob_result);

           use_image = ((absl::GetFlag(FLAGS_filtered_only)
                             ? blob_result.filtered_blobs.size()
                             : blob_result.unfiltered_blobs.size()) > 0);
         }

         if (use_image) {
           if (!absl::GetFlag(FLAGS_image_save_prefix).empty()) {
             std::stringstream image_name;
             image_name << absl::GetFlag(FLAGS_image_save_prefix);

             if (absl::GetFlag(FLAGS_match_timestamps)) {
               // Add the time since match start into the image for debugging.
               // We can match images with the game recording.
               image_name << "match_"
                          << std::chrono::duration_cast<std::chrono::seconds>(
                                 event_loop->monotonic_now() - match_start)
                                 .count()
                          << 's';
             } else {
               image_name << nonmatch_image_count++;
             }

             // Add superstructure state to the filename
             if (superstructure_state !=
                 superstructure::SuperstructureState::IDLE) {
               std::string superstructure_state_name =
                   superstructure::EnumNameSuperstructureState(
                       superstructure_state);
               std::transform(superstructure_state_name.begin(),
                              superstructure_state_name.end(),
                              superstructure_state_name.begin(),
                              [](char c) { return std::tolower(c); });
               image_name << '_' << superstructure_state_name;
             }

             if (robot_moving) {
               image_name << "_moving";
             }

             image_name << ".png";

             cv::imwrite(image_name.str(), image_mat);
           }
           if (absl::GetFlag(FLAGS_display)) {
             cv::imshow("Display", image_mat);
             cv::waitKey(absl::GetFlag(FLAGS_detected_only) ||
                                 absl::GetFlag(FLAGS_filtered_only)
                             ? 10
                             : 1);
           }
         }
       });

   reader.event_loop_factory()->Run();
 }

 void ViewerMain() {
   ReplayData data;
   ReplayRoborio(&data);
   ReplayPi(data);
 }

 }  // namespace
 }  // namespace y2022::vision

 // Quick and lightweight viewer for image logs
 int main(int argc, char **argv) {
   aos::InitGoogle(&argc, &argv);
   y2022::vision::ViewerMain();
 }
	#include "opencv2/calib3d.hpp"
	#include "opencv2/features2d.hpp"
	#include "opencv2/highgui/highgui.hpp"
	#include "opencv2/imgproc.hpp"

	#include "aos/events/logging/log_reader.h"
	#include "aos/events/simulated_event_loop.h"
	#include "aos/init.h"
	#include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
	#include "frc971/input/joystick_state_generated.h"
	#include "frc971/vision/vision_generated.h"
	#include "y2022/control_loops/superstructure/superstructure_status_generated.h"
	#include "y2022/vision/blob_detector.h"

	ABSL_FLAG(std::string, pi, "pi3", "Node name to replay.");
	ABSL_FLAG(std::string, image_save_prefix, "/tmp/img",
	"Prefix to use for saving images from the logfile.");
	ABSL_FLAG(bool, display, false, "If true, display the images with a timeout.");
	ABSL_FLAG(bool, detected_only, false,
	"If true, only write images which had blobs (unfiltered) detected");
	ABSL_FLAG(bool, filtered_only, false,
	"If true, only write images which had blobs (filtered) detected");
	ABSL_FLAG(bool, match_timestamps, false,
	"If true, name the files based on the time since the robot was "
	"enabled (match start). Only consider images during this time");
	ABSL_FLAG(std::string, logger_pi_log, "/tmp/logger_pi/",
	"Path to logger pi log");
	ABSL_FLAG(std::string, roborio_log, "/tmp/roborio/", "Path to roborio log");

	namespace y2022::vision {
	namespace {

	using aos::monotonic_clock;
	namespace superstructure = control_loops::superstructure;

	// Information to extract from the roborio log
	struct ReplayData {
	monotonic_clock::time_point match_start;
	monotonic_clock::time_point match_end;
	std::map<monotonic_clock::time_point, bool> robot_moving;
	std::map<monotonic_clock::time_point, superstructure::SuperstructureState>
	superstructure_states;
	};

	// Extract the useful data from the roborio log to be used for naming images
	void ReplayRoborio(ReplayData *data) {
	data->match_start = monotonic_clock::min_time;
	data->match_end = monotonic_clock::min_time;

	// Open logfiles
	aos::logger::LogReader reader(aos::logger::SortParts(
	aos::logger::FindLogs(absl::GetFlag(FLAGS_roborio_log))));
	reader.Register();
	const aos::Node *roborio =
	aos::configuration::GetNode(reader.configuration(), "roborio");

	std::unique_ptr<aos::EventLoop> event_loop =
	reader.event_loop_factory()->MakeEventLoop("roborio", roborio);

	auto joystick_state_fetcher =
	event_loop->MakeFetcher<aos::JoystickState>("/roborio/aos");
	auto drivetrain_status_fetcher =
	event_loop->MakeFetcher<frc971::control_loops::drivetrain::Status>(
	"/drivetrain");
	auto superstructure_status_fetcher =
	event_loop->MakeFetcher<superstructure::Status>("/superstructure");

	// Periodically check if the robot state updated
	event_loop->AddPhasedLoop(
	[&](int) {
	// Find the match start and end times
	if (joystick_state_fetcher.Fetch()) {
	if (data->match_start == monotonic_clock::min_time &&
	joystick_state_fetcher->enabled()) {
	data->match_start =
	joystick_state_fetcher.context().monotonic_event_time;
	} else {
	if (data->match_end == monotonic_clock::min_time &&
	data->match_start != monotonic_clock::min_time &&
	!joystick_state_fetcher->autonomous() &&
	!joystick_state_fetcher->enabled()) {
	data->match_end =
	joystick_state_fetcher.context().monotonic_event_time;
	}
	}
	}

	// Add whether the robot was moving at a non-negligible speed to
	// the image name for debugging.
	drivetrain_status_fetcher.Fetch();
	if (drivetrain_status_fetcher.get()) {
	// If the robot speed was atleast this (m/s), it is
	// considered moving.
	constexpr double kMinMovingRobotSpeed = 0.5;
	data->robot_moving[drivetrain_status_fetcher.context()
	.monotonic_event_time] =
	(std::abs(drivetrain_status_fetcher->robot_speed()) >=
	kMinMovingRobotSpeed);
	}

	superstructure_status_fetcher.Fetch();
	if (superstructure_status_fetcher.get()) {
	data->superstructure_states[superstructure_status_fetcher.context()
	.monotonic_event_time] =
	superstructure_status_fetcher->state();
	}
	},
	std::chrono::milliseconds(50));
	reader.event_loop_factory()->Run();
	}

	template <typename T>
	T ClosestElement(const std::map<monotonic_clock::time_point, T> &map,
	monotonic_clock::time_point now) {
	T closest;

	// The closest element is either the one right above it, or the element before
	// that one
	auto closest_it = map.lower_bound(now);
	// Handle the case where now is greater than all times in the map
	if (closest_it == map.cend()) {
	closest_it--;
	closest = closest_it->second;
	} else {
	// Start off with the closest as the first after now
	closest = closest_it->second;
	const monotonic_clock::duration after_duration = closest_it->first - now;
	closest_it--;

	// If there is a time before, check if that's closer to now
	if (closest_it != map.cbegin()) {
	const monotonic_clock::duration before_duration = now - closest_it->first;
	if (before_duration < after_duration) {
	closest = closest_it->second;
	}
	}
	}

	return closest;
	}

	// Extract images from the pi logs
	void ReplayPi(const ReplayData &data) {
	if (absl::GetFlag(FLAGS_match_timestamps)) {
	CHECK_NE(data.match_start, monotonic_clock::min_time)
	<< "Can't use match timestamps if match never started";
	CHECK_NE(data.match_end, monotonic_clock::min_time)
	<< "Can't use match timestamps if match never ended";
	}

	// Open logfiles
	aos::logger::LogReader reader(aos::logger::SortParts(
	aos::logger::FindLogs(absl::GetFlag(FLAGS_logger_pi_log))));
	reader.Register();
	const aos::Node *pi = aos::configuration::GetNode(reader.configuration(),
	absl::GetFlag(FLAGS_pi));

	std::unique_ptr<aos::EventLoop> event_loop =
	reader.event_loop_factory()->MakeEventLoop("player", pi);

	LOG(INFO) << "Match start: " << data.match_start
	<< ", match end: " << data.match_end;

	size_t nonmatch_image_count = 0;

	event_loop->MakeWatcher(
	"/camera/decimated", [&](const frc971::vision::CameraImage &image) {
	const auto match_start = data.match_start;
	// Find the closest robot moving and superstructure state to now
	const bool robot_moving =
	ClosestElement(data.robot_moving, event_loop->monotonic_now());
	const auto superstructure_state = ClosestElement(
	data.superstructure_states, event_loop->monotonic_now());

	if (absl::GetFlag(FLAGS_match_timestamps)) {
	if (event_loop->monotonic_now() < data.match_start) {
	// Ignore prematch images if we only care about ones during the
	// match
	return;
	} else if (event_loop->monotonic_now() >= data.match_end) {
	// We're done if the match is over and we only wanted match images
	reader.event_loop_factory()->Exit();
	return;
	}
	}

	// Create color image:
	cv::Mat image_color_mat(cv::Size(image.cols(), image.rows()), CV_8UC2,
	(void *)image.data()->data());
	cv::Mat image_mat(cv::Size(image.cols(), image.rows()), CV_8UC3);
	cv::cvtColor(image_color_mat, image_mat, cv::COLOR_YUV2BGR_YUYV);

	bool use_image = true;
	if (absl::GetFlag(FLAGS_detected_only) \|\|
	absl::GetFlag(FLAGS_filtered_only)) {
	// TODO(milind): if adding target estimation here in the future,
	// undistortion is needed
	BlobDetector::BlobResult blob_result;
	BlobDetector::ExtractBlobs(image_mat, &blob_result);

	use_image = ((absl::GetFlag(FLAGS_filtered_only)
	? blob_result.filtered_blobs.size()
	: blob_result.unfiltered_blobs.size()) > 0);
	}

	if (use_image) {
	if (!absl::GetFlag(FLAGS_image_save_prefix).empty()) {
	std::stringstream image_name;
	image_name << absl::GetFlag(FLAGS_image_save_prefix);

	if (absl::GetFlag(FLAGS_match_timestamps)) {
	// Add the time since match start into the image for debugging.
	// We can match images with the game recording.
	image_name << "match_"
	<< std::chrono::duration_cast<std::chrono::seconds>(
	event_loop->monotonic_now() - match_start)
	.count()
	<< 's';
	} else {
	image_name << nonmatch_image_count++;
	}

	// Add superstructure state to the filename
	if (superstructure_state !=
	superstructure::SuperstructureState::IDLE) {
	std::string superstructure_state_name =
	superstructure::EnumNameSuperstructureState(
	superstructure_state);
	std::transform(superstructure_state_name.begin(),
	superstructure_state_name.end(),
	superstructure_state_name.begin(),
	[](char c) { return std::tolower(c); });
	image_name << '_' << superstructure_state_name;
	}

	if (robot_moving) {
	image_name << "_moving";
	}

	image_name << ".png";

	cv::imwrite(image_name.str(), image_mat);
	}
	if (absl::GetFlag(FLAGS_display)) {
	cv::imshow("Display", image_mat);
	cv::waitKey(absl::GetFlag(FLAGS_detected_only) \|\|
	absl::GetFlag(FLAGS_filtered_only)
	? 10
	: 1);
	}
	}
	});

	reader.event_loop_factory()->Run();
	}

	void ViewerMain() {
	ReplayData data;
	ReplayRoborio(&data);
	ReplayPi(data);
	}

	} // namespace
	} // namespace y2022::vision

	// Quick and lightweight viewer for image logs
	int main(int argc, char **argv) {
	aos::InitGoogle(&argc, &argv);
	y2022::vision::ViewerMain();
	}