Source code for simulation.utils.machine_learning.data.rosbag_to_video

import argparse
import os
import shutil

from simulation.utils.machine_learning.data.images_to_video import images_to_video
from simulation.utils.machine_learning.data.rosbag_to_images import rosbag_to_images


[docs]def rosbag_to_video(rosbag_dir: str, output_dir: str, image_topic: str): os.makedirs(output_dir, exist_ok=True) for root, dirs, files in os.walk(rosbag_dir): for name in files: if not name.lower().endswith(".bag"): continue input_file_path = os.path.join(root, name) rosbag_to_images(input_file_path, os.path.join(rosbag_dir, "tmp"), image_topic) output_file_path = os.path.join(output_dir, name.replace(".bag", ".mp4")) images_to_video( os.path.abspath(os.path.join(rosbag_dir, "tmp/*.png")), output_file_path, use_glob=True, ) shutil.rmtree(os.path.join(rosbag_dir, "tmp"))
if __name__ == "__main__": parser = argparse.ArgumentParser(description="Rosbags to Videos") parser.add_argument( "--rosbag_dir", type=str, required=True, help="directory of all rosbags", ) parser.add_argument( "--output_dir", type=str, required=True, help="the output directory for all videos", ) parser.add_argument("--image_topic", default="/camera/image_raw", help="Image topic.") args = parser.parse_args() rosbag_to_video(args.rosbag_dir, args.output_dir, args.image_topic)