diff --git a/.gitignore b/.gitignore index b6e4761..02e203a 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,6 @@ dmypy.json # Pyre type checker .pyre/ + +# Ignore output folder +output/ diff --git a/README.md b/README.md index 174d96b..4a1b2d6 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,7 @@ Use the following steps to extract data from a rosbag, remove image distortions, # Contact * Aman Saraf | [amansaraf99@gmail.com](mailto:amansaraf99@gmail.com) +* Chris Lai | [chris.lai@berkeley.edu](mailto:chris.lai@berkeley.edu) diff --git a/multi_extract.sh b/multi_extract.sh index dd51fab..fa1ff32 100755 --- a/multi_extract.sh +++ b/multi_extract.sh @@ -3,12 +3,14 @@ # ---------------------------------------------------------------------------- # # DEFAULT VARIABLES # # ---------------------------------------------------------------------------- # -DATA_DIR_DEFAULT="/media/Public/ROSBAG_BACKUPS/rosbag2_2022_09_21-12_58_49" -VERBOSE_DEFAULT=1 +DATA_DIR_DEFAULT="/media/art-berk/DRIVE2_ART/rosbags/" +VERBOSE_DEFAULT=0 UNDISTORT_DEFAULT=1 -CALIB_DIR_DEFAULT="/home/roar/ART/perception/Camera/Calibration_new/" -OUTPUT_BASE_DIR_DEFAULT="/media/Public/Lucas_Oil/" +# CALIB_DIR_DEFAULT="/home/roar/ART/perception/Camera/Calibration_new/" +CALIB_DIR_DEFAULT="/home/art-berk/IAC_dataset_maker/putnam_calib/" +OUTPUT_BASE_DIR_DEFAULT="/home/art-berk/IAC_dataset_maker/output/" MAKE_VID_DEFAULT=1 +USE_COMPRESSED_DEFAULT=0 # ---------------------------------------------------------------------------- # # PARSE ENVIRONMENT VARIABLES # @@ -52,6 +54,18 @@ else echo "----------------------------------------------------------------------------" fi +# ------------------------ USE_COMPRESSED SETTINGS --------------------------- # +if [ -z ${USE_COMPRESSED+x} ]; then + echo "----------------------------------------------------------------------------" + USE_COMPRESSED=$USE_COMPRESSED_DEFAULT + echo " USE_COMPRESSED not defined. Setting USE_COMPRESSED to $USE_COMPRESSED. " + echo "----------------------------------------------------------------------------" +else + echo "----------------------------------------------------------------------------" + echo " USE_COMPRESSED has been defined as $USE_COMPRESSED " + echo "----------------------------------------------------------------------------" +fi + # --------------------------- UNDISTORTION SETTINGS -------------------------- # if [ -z ${UNDISTORT:+x} ]; then echo "----------------------------------------------------------------------------" @@ -90,8 +104,8 @@ else fi # --------------------- ENVIRONMENT VARIABLE PARSING END --------------------- # -echo "\nPausing for 10 seconds. Press Ctrl+C to quit if the above settings are not correct.\n" -for i in 1 2 3 4 5 6 7 8 9 10 +echo "\nPausing for 5 seconds. Press Ctrl+C to quit if the above settings are not correct.\n" +for i in 1 2 3 4 5 do echo "$i seconds passed" sleep 1s @@ -102,9 +116,20 @@ done # Find all rosbags at datadir and extract data # # ---------------------------------------------------------------------------- # echo "\n\nSearching DATA_DIR for ROSBAGS now...\n" -sleep 5s -find "$DATA_DIR" -iname "*.db3" -print0 | xargs -0 -I file dirname file | sort | uniq | while read d; do - ROSBAG_NAME=$(basename "$d") +root_dir=$(pwd) +sleep 2s +# find "$DATA_DIR" \( -iname "*.db3" -o -iname "*.mcap" \) -print0 | xargs -0 -I file dirname file | while read d; do +# find "$DATA_DIR" \( -iname "*.db3" -o -iname "*.mcap" \) -print0 | xargs -0 -I file dirname file | sort -u | while IFS= read -r d; do +# echo "$d" +# done +find "$DATA_DIR" \( -iname "*.db3" -o -iname "*.mcap" \) -print0 | xargs -0 -I file dirname file | sort -u > /tmp/tempfile.txt + + +# find "$DATA_DIR" \( -iname "*.db3" -o -iname "*.mcap" \) -print0 | xargs -0 -I file dirname file | sort -u | while IFS= read -r d; do +while IFS= read -r line; do + + ROSBAG_NAME=$(basename "$line") + # echo "ROSBAG_NAME $ROSBAG_NAME" # ---------------------------------------------------------------------------- # # NOTE: SPECIFY OUTPUT DIR # @@ -115,35 +140,63 @@ find "$DATA_DIR" -iname "*.db3" -print0 | xargs -0 -I file dirname file | sort | # ------------------------------ Debug Verbosity ----------------------------- # echo "----------------------------------------------------------------------------" - echo "Found ROSBAG: $d" + echo "Found ROSBAG: $line" echo "Extracting images to: $OUTPUT_DIR" echo "----------------------------------------------------------------------------" # -------------------------- Create output Directory ------------------------- # - mkdir -p $OUTPUT_DIR + mkdir -p "$OUTPUT_DIR" # ------------------ Begin Extraction Based on User Setting ------------------ # + cd "$root_dir" + if [ $VERBOSE -eq 1 ]; then if [ $UNDISTORT -eq 1 ]; then - python3 ros2bag_image_extractor.py "$d" $OUTPUT_DIR -vup $CALIB_DIR + if [ $USE_COMPRESSED -eq 1 ]; then + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -vucp "$CALIB_DIR" + else + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -vup "$CALIB_DIR" + fi else - python3 ros2bag_image_extractor.py "$d" $OUTPUT_DIR -v + if [ $USE_COMPRESSED -eq 1 ]; then + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -vc + else + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -v + fi fi else if [ $UNDISTORT -eq 1 ]; then - python3 ros2bag_image_extractor.py "$d" $OUTPUT_DIR -up $CALIB_DIR + if [ $USE_COMPRESSED -eq 1 ]; then + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -ucp "$CALIB_DIR" + else + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -up "$CALIB_DIR" + fi else - python3 ros2bag_image_extractor.py "$d" $OUTPUT_DIR + if [ $USE_COMPRESSED -eq 1 ]; then + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" -c + else + python3 ros2bag_image_extractor.py "$line" "$OUTPUT_DIR" + fi fi fi # --------------------- Convert Extracted Images to Video -------------------- # - if [ $MAKE_VID -eq 1 ]; then - for camera_output_dir in $OUTPUT_DIR/*/; do - cd $camera_output_dir - base_name=$(basename "$camera_output_dir") - ffmpeg -framerate 50 -pattern_type glob -i '*.jpg' -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p ../$base_name.mp4 - done - fi + if [ ! -d "$OUTPUT_DIR" ] || [ -z "$(ls -A $OUTPUT_DIR)" ]; then # Check if the output directory is not empty first + if [ $MAKE_VID -eq 1 ]; then + for camera_output_dir in "$OUTPUT_DIR"/*/; do + base_name=$(basename "$camera_output_dir") -done + # Ah, that narrows things down a bit. The behavior you're seeing could be related to ffmpeg inadvertently reading from standard input, which affects the subsequent iterations of the loop that reads from /tmp/tempfile.txt. + # Here's a solution: redirect the standard input of ffmpeg to /dev/null. This ensures that ffmpeg won't interfere with the input being read by the loop. + # Modify the ffmpeg line as follows: + ffmpeg -framerate 50 -pattern_type glob -i "$camera_output_dir/*.jpg" -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p "./output/$ROSBAG_NAME/$base_name.mp4" > /dev/null 2>&1 < /dev/null + + done + fi + fi + sleep 1s +# done +done < /tmp/tempfile.txt + +rm /tmp/tempfile.txt +find $OUTPUT_BASE_DIR_DEFAULT -type d -empty -delete # Cleanup Empty Directories \ No newline at end of file diff --git a/organizer.py b/organizer.py new file mode 100644 index 0000000..c88e0f1 --- /dev/null +++ b/organizer.py @@ -0,0 +1,27 @@ +import os +import shutil + +# Define the paths +first_path = '/home/art-berk/IAC_dataset_maker/output' +second_path = '/media/art-berk/DRIVE2_ART/rosbags' + +# Ensure 'rosbag_with_images_extracted' directory exists +destination_folder = os.path.join(second_path, 'rosbag_with_images_extracted') +os.makedirs(destination_folder, exist_ok=True) + +# Get a set of all directory names in first path +first_path_dirs = {dir_name for dir_name in os.listdir(first_path) + if os.path.isdir(os.path.join(first_path, dir_name))} + +# Iterate over all directories in the second path +for dir_name in os.listdir(second_path): + dir_path = os.path.join(second_path, dir_name) + # Check if it's a directory and not the destination directory + if dir_name not in first_path_dirs or not os.path.isdir(dir_path) or dir_path == destination_folder: + continue # Skip if no match or it's not a directory or if it is the destination directory + + # Move matched directory to the 'rosbag_with_images_extracted' directory + shutil.move(dir_path, destination_folder) + print(f"Moved '{dir_name}' to '{destination_folder}'") + +print("Operation completed.") diff --git a/putnam_calib/vimba_front_left.yaml b/putnam_calib/vimba_front_left.yaml new file mode 100644 index 0000000..8b9a52c --- /dev/null +++ b/putnam_calib/vimba_front_left.yaml @@ -0,0 +1,23 @@ +camera_matrix: + rows: 3 + cols: 3 + data: + - 243.762740 + - 0.000000 + - 259.440156 + - 0.000000 + - 244.236210 + - 191.623041 + - 0.000000 + - 0.000000 + - 1.000000 + +distortion_coefficients: + rows: 1 + cols: 5 + data: + - -0.179154 + - 0.041029 + - -0.001240 + - 0.000456 + - 0.000000 \ No newline at end of file diff --git a/putnam_calib/vimba_front_left_center.yaml b/putnam_calib/vimba_front_left_center.yaml new file mode 100644 index 0000000..60c4fcb --- /dev/null +++ b/putnam_calib/vimba_front_left_center.yaml @@ -0,0 +1,23 @@ +camera_matrix: + rows: 3 + cols: 3 + data: + - 1637.855960 + - 0.000000 + - 546.820824 + - 0.000000 + - 1629.665877 + - 367.335716 + - 0.000000 + - 0.000000 + - 1.000000 + +distortion_coefficients: + rows: 1 + cols: 5 + data: + - -0.316097 + - 0.461574 + - -0.001655 + - 0.004033 + - 0.000000 diff --git a/putnam_calib/vimba_front_right.yaml b/putnam_calib/vimba_front_right.yaml new file mode 100644 index 0000000..a7c7aad --- /dev/null +++ b/putnam_calib/vimba_front_right.yaml @@ -0,0 +1,23 @@ +camera_matrix: + rows: 3 + cols: 3 + data: + - 241.574152 + - 0.000000 + - 245.622459 + - 0.000000 + - 240.962419 + - 188.186666 + - 0.000000 + - 0.000000 + - 1.000000 + +distortion_coefficients: + rows: 1 + cols: 5 + data: + - -0.182866 + - 0.042365 + - -0.001724 + - -0.000473 + - 0.000000 diff --git a/putnam_calib/vimba_front_right_center.yaml b/putnam_calib/vimba_front_right_center.yaml new file mode 100644 index 0000000..2285b26 --- /dev/null +++ b/putnam_calib/vimba_front_right_center.yaml @@ -0,0 +1,23 @@ +camera_matrix: + rows: 3 + cols: 3 + data: + - 1669.803681 + - 0.000000 + - 550.637949 + - 0.000000 + - 1667.774209 + - 261.517346 + - 0.000000 + - 0.000000 + - 1.000000 + +distortion_coefficients: + rows: 1 + cols: 5 + data: + - -0.285679 + - -0.003097 + - -0.007689 + - -0.000588 + - 0.000000 \ No newline at end of file diff --git a/ros2bag_image_extractor.py b/ros2bag_image_extractor.py index 94693fb..13117a4 100644 --- a/ros2bag_image_extractor.py +++ b/ros2bag_image_extractor.py @@ -13,9 +13,12 @@ import cv2 # --------------------------- ROS related Libraries -------------------------- # -from rosbags.rosbag2 import Reader -from rosbags.serde import deserialize_cdr from cv_bridge import CvBridge +import rosbag2_py + +from rclpy.serialization import deserialize_message +from rosidl_runtime_py.utilities import get_message + # ---------------------------------------------------------------------------- # @@ -25,15 +28,19 @@ # ------------------------ Check if directory is valid ----------------------- # def dir_path(string): if os.path.isdir(string): + # print(f"Is a directory! {string}") return string else: + # print(f"Is not a directory! {string}") raise NotADirectoryError(string) # ----------------------- Check if file path is valid ----------------------- # def file_path(string): if os.path.isfile(string): + # print(f"Is a file! {string}") return string else: + # print(f"Is not a string! {string}") raise FileNotFoundError(string) # ------------------------------ Undistort Image ----------------------------- # @@ -58,6 +65,7 @@ def undistort(input, distortion_data): arg_parser.add_argument('rosbag_file_path', help='Path to rosbag to extract the data from', type=dir_path) arg_parser.add_argument('output_dir', help='Path to directory where extracted data should be stored', type=dir_path) arg_parser.add_argument('-u', "--undistort", action="store_true") +arg_parser.add_argument('-c', "--compressed", action="store_false") arg_parser.add_argument('-p', '--camera_info_path', help="Path to folder containing yaml config files for camera info for all cameras", type=dir_path) arg_parser.add_argument('-v', "--verbose", action="store_true") @@ -65,7 +73,16 @@ def undistort(input, distortion_data): args = arg_parser.parse_args() OUTPUT_DIR = args.output_dir -ROSBAG_FILE = args.rosbag_file_path + +# Check if output directory exists +if os.path.exists(OUTPUT_DIR): + # If it does exist, check if the directory is empty. If it is, just leave it. If not, just skip the rosbag. + if os.listdir(OUTPUT_DIR): + print("[script] Directory Exists and is Not Empty! Exiting...") + exit() + +print("[debug] Output Directory: ", OUTPUT_DIR) +ROSBAG_FILE_PATH = args.rosbag_file_path if args.undistort: dir_path(args.camera_info_path) @@ -79,75 +96,144 @@ def undistort(input, distortion_data): # ---------------------------------------------------------------------------- # # ---------------- Create reader instance and open for reading --------------- # -with Reader(ROSBAG_FILE) as reader: +# with Reader(ROSBAG_FILE) as reader: +# Check if the extension is a db3 or mcap +files = os.listdir(ROSBAG_FILE_PATH) +print(f"[script] ROSBAG filepath: {ROSBAG_FILE_PATH}") +for file in files: + if file.endswith(".db3"): + store_type = "sqlite3" + print("[script] Detected Input bag is a db3 file.") + + elif file.endswith(".mcap"): + store_type = "mcap" + print("[script] Detected Input bag is a mcap file.") +if not store_type: + print(f"[script] FATAL ERROR: Input bag is not a db3 or mcap file") + exit() + +reader = rosbag2_py.SequentialReader() + +# ----------------------------- OBTAIN ALL TOPICS ---------------------------- # +# Opens the bag files and sets the converter options +try: + reader.open( + rosbag2_py.StorageOptions(uri=ROSBAG_FILE_PATH, storage_id=store_type), + rosbag2_py.ConverterOptions( + input_serialization_format="cdr", output_serialization_format="cdr" + ), + ) +except Exception as e: + print(e) + exit() + + +# Check if there are images in the ROSBAG, if not, skip! +# Get all the topic types and +if args.compressed: + + image_topics = { + '/vimba_front_left_center/image/compressed', + '/vimba_front_right_center/image/compressed', + '/vimba_front_left/image/compressed', + '/vimba_front_right/image/compressed', + '/vimba_rear_left/image/compressed', + '/vimba_rear_right/image/compressed', + '/vimba_rear_left/image' , + '/vimba_rear_right/image' , + '/vimba_front_left/image' , + '/vimba_front_left_center/image' , + '/vimba_front_right_center/image' , + '/vimba_front_right/image' + } +else: + image_topics = { + '/vimba_rear_left/image' , + '/vimba_rear_right/image' , + '/vimba_front_left/image' , + '/vimba_front_left_center/image' , + '/vimba_front_right_center/image' , + '/vimba_front_right/image' + } - camera_found = False - for connection in reader.connections: - if args.verbose: - print(connection.topic," : ", connection.msgtype) - if 'camera' in connection.topic: - camera_found = True +TOPIC_TYPES = reader.get_all_topics_and_types() +TYPE_MAP = {TOPIC_TYPES[i].name: TOPIC_TYPES[i].type for i in range(len(TOPIC_TYPES))} - if camera_found is False: - exit() +iterator = dict() - # Iterator dictionary to go over camera messages - iterator = { - '/camera/rear_left/image/compressed' : 0, - '/camera/rear_right/image/compressed' : 0, - '/camera/front_left/image/compressed' : 0, - '/camera/front_left_center/image/compressed' : 0, - '/camera/front_right_center/image/compressed' : 0, - '/camera/front_right/image/compressed' : 0, - '/vimba_rear_left/image' : 0, - '/vimba_rear_right/image' : 0, - '/vimba_front_left/image' : 0, - '/vimba_front_left_center/image' : 0, - '/vimba_front_right_center/image' : 0, - '/vimba_front_right/image' : 0, - } +# Initialize an iterator based on whether or not the topic is in the rosbag +for t in TYPE_MAP: + if t in image_topics: + iterator[t] = 0 - for connection, timestamp, rawdata in reader.messages(): - if connection.topic in iterator.keys(): - - # Update iterator for this topic - iterator[connection.topic] += 1 - - # Extract message from rosbag - msg = deserialize_cdr(rawdata, connection.msgtype) - output_topic = None - if (connection.msgtype == "sensor_msgs/msg/CompressedImage"): - np_arr = np.frombuffer(msg.data, np.uint8) - cv2_msg = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) - output_topic = connection.topic[7:-16] - else: - # Convert to cv2 image - cv2_msg = bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8') - output_topic = connection.topic[:-5] - - # Create a directory for topic in output dir if it does not exist - output_directory = OUTPUT_DIR + output_topic - if not os.path.exists(output_directory): - if args.verbose: - print("Creating Directory: ", output_directory) - os.mkdir(output_directory) - - output_file_path = output_directory + 'Image' + '_' + '{0:010d}'.format(iterator[connection.topic]) + '_' + str(msg.header.stamp.sec) + '_' + str(msg.header.stamp.nanosec) + '.jpg' - - # Undistort Image before Saving - if args.undistort: - # If distortion parameters not loaded then load them once - if connection.topic[1:-6] not in distortion_dict: - yaml_file_path = args.camera_info_path + connection.topic[1:-6] + '.yaml' - distortion_fp = open(file_path(yaml_file_path)) - distortion_dict[connection.topic[1:-6]] = yaml.safe_load(distortion_fp) - cv2_msg = undistort(cv2_msg, distortion_dict[connection.topic[1:-6]]) +if len(iterator) == 0: + print("[script] No Images to extract from this rosbag. Exiting...") + # If no camera topics are found, close the ROSBAG and return. TODO: Check if ffmpeg is happy about this + del reader + exit() + +counter = 0 - # Save Image +while reader.has_next(): + + # Read the next message + topic_name, data, timestamp = reader.read_next() + + if topic_name in iterator.keys(): + # Update iterator for this topic + iterator[topic_name] += 1 + + # Extract message from rosbag + msg_type = TYPE_MAP[topic_name] + msg_ser = get_message(msg_type) + msg = deserialize_message(data, msg_ser) + output_topic = None + if (args.compressed and msg_type == "sensor_msgs/msg/CompressedImage"): + np_arr = np.frombuffer(msg.data, np.uint8) + cv2_msg = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) + output_topic = topic_name[7:-17] + else: + # Convert to cv2 image + cv2_msg = bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8') + output_topic = topic_name[1:-6] + + # Create a directory for topic in output dir if it does not exist + # print("Output Topic: ", output_topic) + output_directory = os.path.join(OUTPUT_DIR, output_topic) + + if not os.path.exists(output_directory): if args.verbose: - print('Saving ' + output_file_path) + print("Creating Directory: ", output_directory) + os.mkdir(output_directory) + + output_file_path = os.path.join(output_directory, 'Image' + '_' + '{0:010d}'.format(iterator[topic_name]) + '_' + str(msg.header.stamp.sec) + '_' + str(msg.header.stamp.nanosec) + '.jpg') - if not cv2.imwrite(output_file_path, cv2_msg): - raise Exception("Could not write image") - \ No newline at end of file + # Undistort Image before Saving + if args.undistort: + # If distortion parameters not loaded then load them once + if topic_name[1:-6] not in distortion_dict: + yaml_file_path = args.camera_info_path + topic_name[1:-6] + '.yaml' + try: + distortion_fp = open(file_path(yaml_file_path)) + except: + del reader + print("[script] FATAL ERROR: Opening .yaml file failed. NOT A FILE") + exit() + distortion_dict[topic_name[1:-6]] = yaml.safe_load(distortion_fp) + cv2_msg = undistort(cv2_msg, distortion_dict[topic_name[1:-6]]) + + # Save Image + + + if args.verbose: + print('Saving ' + output_file_path) + counter += 1 + if counter %1000 == 0: + print(f"Processed {counter} Images") + + if not cv2.imwrite(output_file_path, cv2_msg): + raise Exception("Could not write image") + +# Close the bag file +del reader