diff --git a/dataset_generation/add_camera_info.py b/dataset_generation/add_camera_info.py deleted file mode 100755 index e0fcbc8..0000000 --- a/dataset_generation/add_camera_info.py +++ /dev/null @@ -1,66 +0,0 @@ -import rospy -import rosbag -import numpy as np -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import Transform, TransformStamped, Vector3, Quaternion -from sensor_msgs.msg import CameraInfo -from std_msgs.msg import Header -import tf.transformations as tfx -from pathlib import Path -import argparse -import shutil - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-b", "--bag", default="") - parser.add_argument("-o", "--out_bag", default="", - help="Output ROS bag. (Default is to _overwrite_ input bag)") - parser.add_argument("-c", "--camera", default="pylon_camera") - args = parser.parse_args() - - b_path = Path(args.bag) - b_in = rosbag.Bag(str(b_path)) - tmp_bag = '/tmp/tmp.bag' - b_out=rosbag.Bag(tmp_bag, 'w') - # b_out = rosbag.Bag(str(Path(args.out_dir) / Path(b_path.stem))+'.bag', 'w') - - ci = CameraInfo() - ci.distortion_model = 'plumb_bob' - - ci.D = 5*[0] - ci.D[0] = -0.1343 - ci.D[1] = -0.0259 - ci.D[2] = 0.0021 - ci.D[3] = 0.0008 - - ci.K = 9*[0] - ci.K[0] = 2813.64 - ci.K[2] = 969.28 - ci.K[4] = 2808.33 - ci.K[5] = 624.05 - ci.K[8] = 1.0 - - ci.R = 9*[0] - ci.R[0] = 1.0 - ci.R[4] = 1.0 - ci.R[8] = 1.0 - - ci.P = 12*[0] - ci.P[0] = 2766.88 - ci.P[2] = 970.50 - ci.P[5] = 2790.28 - ci.P[6] = 625.22 - ci.P[10] = 1.0 - - for topic, msg, t, cnxn_hdr in b_in.read_messages(return_connection_header=True): - if args.camera in topic and 'Image' in str(type(msg)): - ci.header = msg.header - topic_ns = Path(topic).parent - cam_info = topic_ns / Path('camera_info') - b_out.write(str(cam_info), ci, t) - b_out.write(topic, msg, t) - b_in.close() - b_out.close() - - out_bag = args.out_bag if args.out_bag else args.bag - shutil.move(tmp_bag, out_bag) diff --git a/dataset_generation/apply_cal.py b/dataset_generation/apply_cal.py deleted file mode 100644 index 03a4246..0000000 --- a/dataset_generation/apply_cal.py +++ /dev/null @@ -1,69 +0,0 @@ -import rospy -import rosbag -import numpy as np -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import Transform, TransformStamped, Vector3, Quaternion -from std_msgs.msg import Header -import tf.transformations as tfx -from pathlib import Path -import argparse -import shutil -from tqdm import tqdm -from deconflict_tf import fix_conflicts, update_tf_chain - - -def main(args): - tmp_bag = "/tmp/temp.bag" - overwrite = args.out_bag == args.bag or not args.out_bag - b_out = rosbag.Bag(tmp_bag, 'w') if overwrite else rosbag.Bag(args.out_bag, 'w') - - tf_static_msg = None - tf_static_t = None - tf_static_connection_header = None - - for topic, msg, t, cnxn_hdr in b_cal.read_messages(topics=['/tf_static', '/tf']): - if tf_static_msg: - tf_static_msg.transforms.extend(msg.transforms) - else: - tf_static_msg = msg - tf_static_t = t - tf_static_connection_header = cnxn_hdr - b_cal.close() - - - for topic, msg, t in b.read_messages(): - if topic == "/tf_static": - update_tf_chain(static_tfs, msg) - elif topic == "/tf": - update_tf_chain(dynamic_tfs, msg) - viz_tree('/tmp/tree_before.dot',static_tfs, dynamic_tfs) - - - for b in args.bags: - b_in = rosbag.Bag(b) - for topic, msg, t, cnxn_hdr in tqdm(b_in.read_messages(return_connection_header=True)): - if topic == '/tf_static': - if tf_static_msg: - tf_static_msg.transforms.extend(msg.transforms) - else: - tf_static_msg = msg - tf_static_t = t - tf_static_connection_header = cnxn_hdr - else: - b_out.write(topic, msg, t) - b_in.close() - b_out.write('/tf_static', tf_static_msg, tf_static_t, - connection_header=tf_static_connection_header) - b_out.close() - - if overwrite: - shutil.move(tmp_bag, args.bag) - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-b", "--bag", nargs='+', default=[], help="bags to update") - parser.add_argument("-c", "--cal_bag", nargs='+', default=[], help="bag with cal information in tf_static") - parser.add_argument("-o", "--out_bag", default="", - help="Output ROS bag. (Default is to _overwrite_ input bag)") - args = parser.parse_args() - main(args) diff --git a/dataset_generation/bag_extract.sh b/dataset_generation/bag_extract.sh deleted file mode 100755 index a28eb68..0000000 --- a/dataset_generation/bag_extract.sh +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env bash - -OPTIND=1 - -class='testclass' -overwrite=0 -last_n=0 -bag_dir="$HOME/data/rcta_objects/barrier" -sync="approx" -color="/camera/color/image_rect_color" -depth="/camera/aligned_depth_to_color/image_raw" -out="" -undistort=0 -all_topics=() -all_types=() -topics="" -cam_info="" -sensor_name="" -sensor_ns="" -sensor_type="" -usage="$(basename "$0") [-h] [-b bag_dir (default '$bag_dir')] [-c class (default '$class')] -w overwrite sequence dirs, starting at seq 0. Use with caution. (no argument for this flag) [-l last_n (default 0. A value > 0 will process only the most recent n bag files)] [-s sync (default 'approx')] [-u undistort. No argument for this flag] [-o output_dir (default /../)] [-t topics. Comma separated list of topics]" - -while getopts "ho:wua:b:s:i:c:l:t:" opt; do # if a colon follows an option, the option comes with a argument - case "$opt" in - h) - echo $usage - exit 0 - ;; - o) - out=$OPTARG - ;; - w) - overwrite=1 - ;; - u) - undistort=1 - ;; - c) - class=$OPTARG - ;; - b) - bag_dir=$OPTARG - ;; - s) - sync=$OPTARG - ;; - i) - info=$OPTARG - ;; - t) - IFS=',' topics=($OPTARG) - ;; - l) - last_n=$OPTARG - esac -done -shift $((OPTIND-1)) -#[ "$1" = "--" ] && shift - -if [[ -z $out ]]; then - out="${bag_dir}/../" -fi - -synchronize() { - b=$1 - if [[ $sync == "approx" ]]; then - ./synchronize_frames.py --dataset $b --topic_filter ${all_topics[@]} --type_filter ${all_types[@]} --approx - else - ./synchronize_frames.py --dataset $b --topic_filter ${all_topics[@]} --type_filter ${all_types[@]} - fi -} - -get_sensor_ns() { - t=$1 - topic_stem=$(echo $t | sed 's|.*/||') - sensor_ns=$(echo $t | sed 's|\(.*\)/.*|\1|') - ns_stem=$(echo $sensor_ns | sed 's|.*/||') - if [[ $topic_stem == "compressed" ]]; then - sensor_ns=$(echo $sensor_ns | sed 's|\(.*\)/.*|\1|') - elif [[ $ns_stem == "color" ]]; then - sensor_ns=$(echo $sensor_ns | sed 's|\(.*\)/.*|\1|') - fi - sensor_name=$(echo $sensor_ns | sed 's|.*/||') -} - -get_cam_info() { - b=$1 - t=$2 - cam_info="" - get_sensor_ns $t - ns_topics=$(rosbag info $b | grep $sensor_ns) - info=$(echo $ns_topics | grep "sensor_msgs/CameraInfo") - if [[ $info ]]; then - cam_info=$(echo $info | tr -s ' ' | cut -c2- | cut -d ' ' -f1) - fi -} - -get_topic_type() { - b=$1 - t=$2 - bag_info=$(rosbag info $b | grep $t) - if [[ $(echo $bag_info | grep "sensor_msgs/Image") ]]; then - if [[ $(echo $t | grep "depth") ]]; then - sensor_type="DepthImage" - else - sensor_type="Image" - fi - elif [[ $(echo $bag_info | grep "sensor_msgs/CameraInfo") ]]; then - sensor_type="CameraInfo" - elif [[ $(echo $bag_info | grep "sensor_msgs/PointCloud2") ]]; then - sensor_type="PointCloud2" - else - echo "Unknown type for topic $t" - fi - -} - -get_topic_types() { - b=$1 - for t in ${topics[@]}; do - get_topic_type $b $t - if [[ $sensor_type == "Image" || $sensor_type == "DepthImage" ]]; then - all_topics+=($t) - all_types+=("sensor_msgs/Image") - get_cam_info $b $t - if [[ -n $cam_info ]]; then - all_topics+=($cam_info) - all_types+=("sensor_msgs/CameraInfo") - else - echo "Unable to find CameraInfo for $t" - fi - elif [[ $sensor_type == "PointCloud2" ]]; then - all_topics+=($t) - all_types+=("sensor_msgs/PointCloud2") - fi - done -} - -extract() { - b=$1 - i=$2 - for t in ${all_topics[@]}; do - get_topic_type $b $t - get_sensor_ns $t - if [[ $sensor_type == "DepthImage" ]]; then - ./extract_images.py --bag_file $b --undistort $undistort --output_dir $out/$class/$i/$sensor_name/ --image_topic $t --image_info_topic $cam_info --image_enc mono16 - elif [[ $sensor_type == "Image" ]]; then - ./extract_images.py --bag_file $b --undistort $undistort --output_dir $out/$class/$i/$sensor_name/ --image_topic $t --image_info_topic $cam_info --image_enc bgr8 --file_format jpg - elif [[ $sensor_type == "PointCloud2" ]]; then - ./extract_ply.py --bag_file $b --output_dir $out/$class/$i/$sensor_name --pcd_topic $t - fi - done -} - -if [[ $last_n == 0 ]]; then - readarray -t all_bags<<<"$(ls -hardt $bag_dir/*bag | grep -v synced.bag)" -else - readarray -t all_bags<<<"$(ls -hardt $bag_dir/*bag | grep -v synced.bag | tail -n $last_n)" -fi - -echo "Num bags ${#all_bags[@]}" -if [[ ! -f $out/$class ]]; then - mkdir -p $out/$class -fi - -if [[ $overwrite == 1 ]]; then - seq_id=0 -else - seq_id=`find $out/$class -maxdepth 1 -type d -name "0*" | wc -l` -fi - -echo "seq id $seq_id" - -if [[ ! $(rostopic list) ]]; then - roscore& -fi - -for b in "${all_bags[@]}"; do - echo $b - seq_dir=$(printf "%05d" $seq_id) - echo "seq dir $seq_dir" - bag_synced="$(dirname $b)/$(basename $b ".bag")_synced.bag" - echo "output $out/$class/$seq_dir/rgb" - num_sensors="${#topics[@]}" - get_topic_types $b - echo "All topics ${topics[@]}" - echo "First topic ${topics[0]}" - if [[ $num_sensors -gt 1 ]]; then - echo "syncing $bag_synced" - synchronize $b - extract $bag_synced $seq_dir - else - extract $b $seq_dir - fi - if [[ $(ls $out/$class/$seq_dir/$sensor_name | wc -l) != "0" ]]; then - seq_id=$(($seq_id+1)) - fi -done diff --git a/dataset_generation/construct_dataset.py b/dataset_generation/construct_dataset.py deleted file mode 100644 index aeba90e..0000000 --- a/dataset_generation/construct_dataset.py +++ /dev/null @@ -1,316 +0,0 @@ -#!/usr/bin/env python -from __future__ import print_function -import os -import rospy -import argparse -import urllib.request -import numpy as np -import rosbag -from pathlib import Path -from tqdm import tqdm -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext -import json -import csv -import cv2 -import yaml -# import tf2_ros -# from tf2_msgs.msg import TFMessage - -def hex_to_rgb(h): - h = h.lstrip('#') - hlen = len(h) - return tuple(int(h[i:i+hlen//3], 16) for i in range(0, hlen, hlen//3)) - -def hex_to_bgr(h): - return tuple(reversed(hex_to_rgb(h))) - -def get_ontology(ontology_path, label_override_fname): - id_color_map = {} - id_rgb_map = {} - id_name_map = {} - id_remap = {} - if ontology_path.exists(): - ontology_lines = open(str(ontology_path),'rU').readlines() - for parts in csv.reader(ontology_lines[1:], quotechar='"', delimiter=',', - quoting=csv.QUOTE_ALL, skipinitialspace=True): - color_hex = parts[0] - lab_id = parts[3] - lab_name = parts[2].replace(' ','') - id_color_map[int(lab_id)] = list(hex_to_bgr(color_hex)) - id_rgb_map[int(lab_id)] = list(hex_to_rgb(color_hex)) - id_name_map[int(lab_id)] = lab_name - - if label_override_fname: - with open(label_override_fname) as f: - label_remap = yaml.load(f.read()) - name_id_map = {v:k for k,v in id_name_map.items()} - id_remap = {name_id_map[k]:name_id_map[v] for k,v in label_remap.items()} - for k,v in id_remap.items(): - id_color_map.pop(k) - id_rgb_map.pop(k) - id_name_map = {k:label_remap.get(v,v) for k,v in id_name_map.items()} - - num_ontology_blocks = int(np.sqrt(len(id_color_map.keys()))) - if num_ontology_blocks**2 != len(id_color_map.keys()): - num_ontology_blocks += 1 - block_sz = 200 - ontology_img = np.zeros((block_sz*num_ontology_blocks,block_sz*num_ontology_blocks,3), dtype=np.uint8) - ont_block_idx = 0 - block_txt_map = {} - block_color_map = {} - font_type = cv2.FONT_HERSHEY_SIMPLEX - font_scale = 1 - font_thickness = 1 - for o,c in id_color_map.items(): - j_idx = ont_block_idx // num_ontology_blocks - i_idx = (ont_block_idx - j_idx * num_ontology_blocks) % num_ontology_blocks - ont_start_pt = np.array([block_sz*i_idx,block_sz*j_idx]) - ont_end_pt = ont_start_pt + np.array([block_sz,block_sz]) - ont_text_pt = (ont_start_pt + ont_end_pt - np.array([block_sz,0])) // 2 - ontology_img[ont_start_pt[1]:ont_end_pt[1],ont_start_pt[0]:ont_end_pt[0],:] = c - block_txt_map[tuple(ont_text_pt)] = id_name_map[o]+'('+str(o)+')' - block_color_map[tuple(ont_text_pt)] = c - ont_block_idx += 1 - for p,t in block_txt_map.items(): - font_sz, _ = cv2.getTextSize(t, font_type, font_scale, font_thickness) - fscale = min(font_scale, block_sz / font_sz[0]) - cv2.putText(ontology_img, t, p, font_type, fscale, (0,0,0), font_thickness) - - cv2.imwrite(str(ontology_path.with_suffix('.png')), ontology_img) - - with open(str(ontology_path.with_suffix('.yaml')), 'w') as f: - f.write(yaml.dump([id_name_map, id_rgb_map])) - - # dump_jsonl([id_name_map, id_rgb_map], str(ontology_path.with_suffix('.jsonl'))) - - return (id_color_map, id_name_map, id_remap) - -def dump_jsonl(data, output_path, append=False): - """ - Write list of objects to a JSON lines file. - """ - mode = 'a+' if append else 'w' - with open(output_path, mode, encoding='utf-8') as f: - for line in data: - json_record = json.dumps(line, ensure_ascii=False) - f.write(json_record + '\n') - # print('Wrote {} records to {}'.format(len(data), output_path)) - -def load_jsonl(input_path) -> list: - """ - Read list of objects from a JSON lines file. - """ - data = [] - with open(input_path, 'r', encoding='utf-8') as f: - for line in f: - data.append(json.loads(line.rstrip('\n|\r'))) - # print('Loaded {} records from {}'.format(len(data), input_path)) - return data - -def stamp_str_extract(fname): - stem = Path(fname).stem - stamp_idx = stem.rfind('-') - return stem[stamp_idx+1:] - -def stamp_from_filename(fname): - stem = Path(fname).stem - msec_idx = stem.rfind('_') - msecs = float(stem[msec_idx+1:]) - sec_idx = stem[:msec_idx].rfind('-') - secs = float(stem[sec_idx+1:msec_idx]) - stamp = secs + (msecs / 1000.0) - return stamp - -def seq_in_bag(seq, bag): - bag_start = bag.get_start_time() - bag_end = bag.get_end_time() - print('bag',bag_start,bag_end,'seq',seq[0],seq[1]) - return (seq[0] >= bag_start and seq[1] <= bag_end) - -def stamp_in_seq(stamp, seq): - return (seq[0]<= stamp and seq[1] >= stamp) - -def get_associated_fname(seq, fname, args): - fname_stamp = stamp_str_extract(fname) - datadirs = [d for d in Path(seq).iterdir() if d.is_dir() if 'label_' not in str(d)] - data_files = sorted([df for dd in datadirs for df in dd.iterdir() - if df.suffix in args.data_extensions]) - for df in data_files: - if fname_stamp in str(df): - return (df.parent.stem, df.stem) - return None, None - -def write_color(collection_path, sid, sensor_name, fname_stem, id_color_map, - ext_out=".png", label_img_fmt='rgb'): - label_id_name = collection_path / Path(sid) / Path(str(sensor_name)+'_label_id/'+fname_stem+ext_out) - color_file = collection_path / Path(sid) / Path(str(sensor_name)+'_label_color/'+fname_stem+ext_out) - if not color_file.parent.exists(): - color_file.parent.mkdir() - img = cv2.imread(str(label_id_name), cv2.IMREAD_UNCHANGED) - if img is None: - print('Empty',str(label_id_name)) - return - if len(img.shape) < 3: - img_mono = img - else: - img_mono = img[:,:,2] if label_img_fmt == 'rgb' else img[0,:,:] - img_flat = np.reshape(img_mono, (-1)) - img_bgr_flat = np.array([id_color_map.get(p,(0,0,0)) for p in img_flat], dtype=np.uint8) - img_bgr = img_bgr_flat.reshape(img.shape[0],img.shape[1],3) - cv2.imwrite(str(color_file), img_bgr) - # cv2.imwrite('/tmp/label_color.png', img_bgr) - -def write_label(collection_path, sid, sensor_name, fname_stem, id_remap, - ext_out=".png", label_img_fmt='rgb'): - label_id_name = collection_path / Path(sid) / Path(str(sensor_name)+'_label_id/'+fname_stem+ext_out) - img = cv2.imread(str(label_id_name), cv2.IMREAD_UNCHANGED) - if img is None: - print('Empty',str(label_id_name)) - return - if len(img.shape) < 3: - img_mono = img - else: - img_mono = img[:,:,2] if label_img_fmt == 'rgb' else img[0,:,:] - img_flat = np.reshape(img_mono, (-1)) - img_remapped_flat = np.array([id_remap.get(p,p) for p in img_flat], dtype=img_flat.dtype) - img_remapped = img_remapped_flat.reshape(img.shape[0],img.shape[1],-1) - cv2.imwrite(str(label_id_name), img_remapped) - # cv2.imwrite('/tmp/label_id.png', img_remapped) - -def main(args): - root_path = Path(args.root_dir) - collection_path = root_path / Path(args.collection_name) - - # Get sequence start/end timestamps - seq_stamp_map = {} # sequence id: (start time, end time) - stamp_seq_map = {} - sdirs = [d for d in collection_path.iterdir() if d.is_dir()] - for sd in sdirs: - if args.seq_dirs and str(sd.stem) not in args.seq_dirs: - continue - print(sd) - ddl = [d for d in Path(sd).iterdir() if d.is_dir()] - data_files = sorted([df for dd in ddl for df in dd.iterdir() - if df.suffix in args.data_extensions]) - stamps_sorted = sorted([stamp_from_filename(df) for df in data_files]) - if not stamps_sorted: - continue - timestamp_begin = stamps_sorted[0] - timestamp_end = stamps_sorted[-1] - print('ts_begin',timestamp_begin,'ts_end',timestamp_end) - seq_stamp_map[str(sd)] = (timestamp_begin, timestamp_end) - stamp_seq_map[timestamp_begin] = str(sd) - sorted_seqs = [stamp_seq_map[k] for k in sorted(stamp_seq_map.keys())] - sorted_stamps = [seq_stamp_map[sk] for sk in sorted_seqs] - print('All sequences in order: ',sorted_seqs) - print('All stamps in order', sorted_stamps) - print('Stamp differences', [sorted_stamps[i][0] - sorted_stamps[i-1][1] - for i in range(len(sorted_stamps)) if i>0]) - - # Process Appen jsonlist files - seq_json_map = {} - seq_img_map = {} - seq_ann_map = {} - - for jlf in args.jsonl: - jlines = load_jsonl(jlf) - image_urls = [] - ann_urls = [] - for jl in jlines: - for jmt in jl['results']['judgments']: - if jmt['data']['broken_link'] == 'true': - continue - image_urls.append(jmt['unit_data']['image_url']) - ann_urls.append(json.loads(jmt['data']['annotation'])['url']) - timestamp_jl = stamp_from_filename(image_urls[len(image_urls) // 2]) - for sid, stamps in seq_stamp_map.items(): - if stamp_in_seq(timestamp_jl, stamps): - seq_json_map.setdefault(sid, []).append([jlf]) - seq_img_map.setdefault(sid, []).extend(image_urls) - seq_ann_map.setdefault(sid, []).extend(ann_urls) - - # Process Appen ontology file - ontology_path = collection_path / Path('ontology.csv') - id_color_map, id_name_map, id_remap = get_ontology(ontology_path, args.override_fname) - - # Extract images - if args.download_anns: - for sid, fnames in seq_img_map.items(): - print('Seq:',sid) - for i, fname in tqdm(enumerate(fnames)): - ext = Path(fname).suffix - ext_out = ext if ext != '.jpg' else '.png' - sensor_name, fname_stem = get_associated_fname(sid, fname, args) - if sensor_name: - out_name = collection_path / Path(sid) / Path(str(sensor_name)+'_label_id/'+fname_stem+ext_out) - if not out_name.parent.exists(): - out_name.parent.mkdir() - try: - urllib.request.urlretrieve(seq_ann_map[sid][i], str(out_name)) - except Exception as e: - print (e) - continue - - # Convert label_id, generate label_color - for sd in sdirs: - if args.seq_dirs and str(sd.stem) not in args.seq_dirs: - continue - ddl = [d for d in Path(sd).iterdir() if d.is_dir() and 'label_id' in str(d)] - data_files = sorted([df for dd in ddl for df in dd.iterdir() - if df.suffix in args.data_extensions]) - for df in tqdm(data_files): - par = df.parent.stem - sid = df.parent.parent.stem - sensor_name = par[0:par.find('_label_id')] - if args.override_fname: - write_label(collection_path, sid, sensor_name, df.stem, id_remap, ext_out=".png") - write_color(collection_path, sid, sensor_name, df.stem, id_color_map, ext_out=".png") - # print(df) - # assert(False) - - - # Check for bags in //bags. Associate with sequences - bag_path = collection_path / Path('bags') - bag_seq_map={} - if bag_path.exists(): - print('bag path',bag_path) - bag_names = [b for b in bag_path.iterdir() if b.suffix == '.bag' and '_synced.bag' not in str(b)] - for b in bag_names: - bag = rosbag.Bag(str(b)) - for sid, stamps in seq_stamp_map.items(): - if seq_in_bag(stamps, bag): - bag_seq_map[b.name] = sid # assumes 1-1 bag/seq mapping (any split bags have already been merged) - # bag_seq_map.setdefault(b.name,[]).append([sid]) - - bag_seq_path = collection_path / Path('bag_sequence_map.yaml') - bag_seq_stem_map = {k:Path(v).stem for k,v in bag_seq_map.items()} - with open(str(bag_seq_path), 'w') as f: - f.write(yaml.dump(bag_seq_stem_map)) - - for b,s in bag_seq_map.items(): - seq_map_path = collection_path / Path(s) / Path('source_bag.txt') - seq_map_path.write_text(b) - - - -if __name__ == '__main__': - parser = argparse.ArgumentParser(description="Associate jsonlines files to sequences") - parser.add_argument("-r", "--root_dir", default="", - help="dataset root dir") - parser.add_argument("-n", "--collection_name", default="rellis", - help="collection name") - parser.add_argument("-s", "--seq_dirs", nargs='+',default=[], - help="sequence_dirs to process, if empty, process entire collection") - parser.add_argument("-j", "--jsonl", nargs='+', default=[], - help="Appen formatted JSONlines files.") - parser.add_argument("-e", "--data_extensions", nargs='+',default=['.png','.jpg'], - help="set of supported extensions of raw data files") - parser.add_argument("-o", "--override_fname", default="", - help="yaml formatted file for class label overrides") - parser.add_argument("-p", "--pngs", action="store_true", - help="convert all jpgs to pngs") - parser.add_argument("-d", "--download_anns", action="store_true", - help="try to retrieve all images") - args = parser.parse_args() - main(args) diff --git a/dataset_generation/deconflict_tf.py b/dataset_generation/deconflict_tf.py deleted file mode 100644 index 0d72ffe..0000000 --- a/dataset_generation/deconflict_tf.py +++ /dev/null @@ -1,145 +0,0 @@ -#!/usr/bin/env python - -import rosbag -import argparse -import shutil -import tf.transformations as tfx - -def new_tf(tfs, tform): - return tform.child_frame_id not in tfs - -def conflicting_tf(tfs, tform): - assert(not new_tf(tfs, tform)) - return tform.header.frame_id not in tfs[tform.child_frame_id] - -def update_tf_chain(tfs, msg): - for tform in msg.transforms: - if new_tf(tfs,tform): - tfs[tform.child_frame_id] = [tform.header.frame_id] - elif conflicting_tf(tfs, tform): - tfs[tform.child_frame_id].append(tform.header.frame_id) - -def find_conflicts(tfs, parent_fix_frames, - child_fix_frames, all_child_frames): - for ctf, ptfs in tfs.iteritems(): - if ctf in all_child_frames: - for ptf in ptfs: - if ptf not in parent_fix_frames: - parent_fix_frames.append(ptf) - child_fix_frames.append(ctf) - else: - all_child_frames.add(ctf) - -def find_invert_tfs(conflicts, static_tfs, dynamic_tfs): - for c, p in conflicts: - if p in static_tfs: - for pf in static_tfs[p]: - conflicts.append((p,pf)) - elif p in dynamic_tfs: - for pf in dynamic_tfs[p]: - conflicts.append((p,pf)) - return conflicts - -def fix_conflicts(b, bo, static_tfs, dynamic_tfs): - parent_fix_frames = [] - child_fix_frames = [] - all_child_frames = set() - find_conflicts(dynamic_tfs, parent_fix_frames, - child_fix_frames, all_child_frames) - find_conflicts(static_tfs, parent_fix_frames, - child_fix_frames, all_child_frames) - conflicts = zip(child_fix_frames, parent_fix_frames) - all_conflicts = find_invert_tfs(conflicts, static_tfs, dynamic_tfs) - - print(all_conflicts) - for topic, msg, t, cnxn_hdr in b.read_messages(return_connection_header=True): - if topic == "/tf_static" or topic == "/tf": - for i, tform in enumerate(msg.transforms): - if (tform.child_frame_id, tform.header.frame_id) in all_conflicts: - msg.transforms[i] = flip_tform(tform) - bo.write(topic, msg, t, connection_header=cnxn_hdr) - - -def flip_tform(tform_msg): - ri = tfx.euler_from_quaternion([tform_msg.transform.rotation.x, - tform_msg.transform.rotation.y, - tform_msg.transform.rotation.z, - tform_msg.transform.rotation.w]) - ti = [tform_msg.transform.translation.x, - tform_msg.transform.translation.y, - tform_msg.transform.translation.z] - tform = tfx.compose_matrix(angles=ri, translate=ti) - ro = tfx.quaternion_from_matrix(tfx.inverse_matrix(tform)) - to = tfx.translation_from_matrix(tfx.inverse_matrix(tform)) - tform_msg_out = tform_msg - child_frame = tform_msg.child_frame_id - parent_frame = tform_msg.header.frame_id - tform_msg_out.child_frame_id = parent_frame - tform_msg_out.header.frame_id = child_frame - tform_msg_out.transform.translation.x = to[0] - tform_msg_out.transform.translation.y = to[1] - tform_msg_out.transform.translation.z = to[2] - tform_msg_out.transform.rotation.x = ro[0] - tform_msg_out.transform.rotation.y = ro[1] - tform_msg_out.transform.rotation.z = ro[2] - tform_msg_out.transform.rotation.w = ro[3] - return tform_msg_out - - -def viz_tree(fname, static_tfs, dynamic_tfs): - all_nodes = set(static_tfs.keys()+dynamic_tfs.keys()) - for vs in static_tfs.itervalues(): - for v in vs: - all_nodes.add(v) - for vs in dynamic_tfs.itervalues(): - for v in vs: - all_nodes.add(v) - with open(fname, 'w') as f: - f.write('digraph {\n') - for n in all_nodes: - f.write(' {} [label={}]\n'.format(n, n)) - for k,vs in static_tfs.iteritems(): - for v in vs: - f.write(' {} -> {}\n'.format(v,k)) - for k,vs in dynamic_tfs.iteritems(): - for v in vs: - f.write(' {} -> {}\n'.format(v,k)) - f.write('}') - -def main(args): - b=rosbag.Bag(args.bag_file) - tmp_bag = '/tmp/tmp.bag' - bo=rosbag.Bag(tmp_bag, 'w') - static_tfs = {} - dynamic_tfs = {} - for topic, msg, t in b.read_messages(): - if topic == "/tf_static": - update_tf_chain(static_tfs, msg) - elif topic == "/tf": - update_tf_chain(dynamic_tfs, msg) - viz_tree('/tmp/tree_before.dot',static_tfs, dynamic_tfs) - - fix_conflicts(b, bo, static_tfs, dynamic_tfs) - b.close() - bo.close() - - static_tfs = {} - dynamic_tfs = {} - bo=rosbag.Bag(tmp_bag) - for topic, msg, t in bo.read_messages(): - if topic == "/tf_static": - update_tf_chain(static_tfs, msg) - elif topic == "/tf": - update_tf_chain(dynamic_tfs, msg) - viz_tree('/tmp/tree_after.dot', static_tfs, dynamic_tfs) - - out_bag = args.out_bag if args.out_bag else args.bag_file - shutil.move(tmp_bag, out_bag) - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Fix tf tree conflicts retroactively", - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--bag_file", default=None, help="Input ROS bag.") - parser.add_argument("--out_bag", default=None, help="Output ROS bag. (Default is to _overwrite_ input bag)") - args = parser.parse_args() - main(args) diff --git a/dataset_generation/extract_images.py b/dataset_generation/extract_images.py deleted file mode 100755 index 2fa10a2..0000000 --- a/dataset_generation/extract_images.py +++ /dev/null @@ -1,182 +0,0 @@ -#!/usr/bin/env python -# -*- coding: utf-8 -*- - -# Copyright 2016 Massachusetts Institute of Technology - -"""Extract images from a rosbag. -""" - -from __future__ import print_function -import os -import argparse - -import cv2 -import numpy as np -import rosbag -from sensor_msgs.msg import Image -from cv_bridge import CvBridge, CvBridgeError -from tqdm import tqdm -import json -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext -from pathlib import Path -import yaml - -def main(args): - - print("Extract images from %s on topic %s into %s" % (args.bag_file, - args.image_topic, - args.output_dir)) - - if args.image_info_topic == None: - print("For image topic", args.image_topic) - args.image_info_topic = args.image_topic[:args.image_topic.rfind('/')] + "/camera_info" - print("Info topic", args.image_info_topic) - - bag = rosbag.Bag(args.bag_file, "r") - bridge = CvBridge() - count = 0 - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - size = None - intrinsics = None - distortion = None - R = None - P = None - for topic, msg, t in bag.read_messages(topics=[args.image_info_topic]): - print("camera_info:", msg) - yaml_info = {} - yaml_info['data_type'] = 'image' - yaml_info['frame_id'] = msg.header.frame_id - yaml_info['width'] = msg.width - yaml_info['height'] = msg.height - - fx = msg.K[0] - cx = msg.K[2] - fy = msg.K[4] - cy = msg.K[5] - # distortion = msg.D - # distortion_model = msg.distortion_model - # yaml_info['intrinsic_matrix'] = [fx, 0.0, 0.0, 0.0, fy, 0.0, cx, cy, 1.0] - yaml_info['K'] = list(msg.K) - yaml_info['D'] = list(msg.D) - yaml_info['R'] = list(msg.R) - yaml_info['P'] = list(msg.P) - # size = (msg.width, msg.height) - # intrinsics = np.array(msg.K, dtype=np.float64, - # copy=True).reshape((3, 3)) - # distortion = np.array(msg.D, dtype=np.float64, - # copy=True).reshape((len(msg.D), 1)) - # R = np.array(msg.R, dtype=np.float64, copy=True).reshape((3, 3)) - # P = np.array(msg.P, dtype=np.float64, copy=True).reshape((3, 4)) - fname_dir = Path(args.output_dir) - with open(str(fname_dir.parent / Path(fname_dir.stem+'_info.yaml')), 'w') as f: - f.write(yaml.dump(yaml_info)) - # json.dump(yaml_info, f) - - with open(os.path.join(args.output_dir + "/..", 'camera_info.txt'), 'w') as f: - f.write("{} {} {} {}\n".format(fx, fy, cx, cy)) - break - - issued_warning = False - issued_undistort_msg = False - for topic, msg, t in tqdm(bag.read_messages(topics=[args.image_topic]), total=bag.get_message_count(args.image_topic)): - try: - cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding=args.image_enc) - except CvBridgeError as e: - if not issued_warning: - print ('Exception:', e, 'trying to force image type conversion') - cv_img_raw = bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough') - cv_img_np = np.array(cv_img_raw) - if 'depth' in args.image_topic: - if np.nanmax(cv_img_np) < 1000: - if not issued_warning: - print ('Converting assumed depth image from m to mm, output will be 16uc1') - cv_img_mm = 1000.0 * cv_img_np - cv_img = cv_img_mm.astype(np.uint16) - else: - cv_img = cv_img_np.astype(np.uint16) - issued_warning = True - - if intrinsics is not None and args.undistort: - if not issued_undistort_msg: - print ('Undistorting raw input image') - issued_undistort_msg = True - cv_img_raw = cv_img - cv_img = cv2.undistort(cv_img_raw, intrinsics, distortion) - if np.any(np.isnan(cv_img)): - print("Predicted image has NaNs") - stamp = msg.header.stamp - stamp_ns_str = str(stamp.nsecs).zfill(9) - cv2.imwrite(os.path.join(args.output_dir, - "frame%06i-%s_%s.%s" % (count, str(stamp.secs), - stamp_ns_str[:3], args.file_format)), cv_img) - # print ("Wrote image %i" % count) - - count += 1 - - print("Wrote {} images to {}".format(count,args.output_dir)) - bag.close() - - return - -if __name__ == '__main__': - # Default from xtion RGB: ./extract_images.py --bag_file ~/data/rcta_box.bag --output_dir ~/data/rcta_objects/box/00001/rgb/ --image_topic /rgb/image --image_enc bgr8 --file_format jpg - # Default from xtion Depth: ./extract_images.py --bag_file ~/data/rcta_box.bag --output_dir ~/data/rcta_objects/box/00001/depth/ --image_topic /depth/image --image_enc mono16 - - """Extract a folder of images from a rosbag. - """ - parser = argparse.ArgumentParser(description="Extract images from a ROS bag.", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog='''Example: - RGB: ./extract_images.py --bag_file ~/data/rcta_box.bag --output_dir ~/data/rcta_objects/box/00001/rgb/ --image_topic /rgb/image --image_enc bgr8 --file_format jpg - Depth: ./extract_images.py --bag_file ~/data/rcta_box.bag --output_dir ~/data/rcta_objects/box/00001/depth/ --image_topic /depth/image --image_enc mono16''') - parser.add_argument("--bag_file", default=None, help="Input ROS bag.") - parser.add_argument("--rosbag_folder", default=None, help="Folder of synchronized rosbags to extract") - parser.add_argument("--output_dir", help="Output directory.") - parser.add_argument("--image_topic", help="Image topic.", default=None) - parser.add_argument("--image_info_topic", help="Info topic.", nargs='?', default=None) - camera_type_param = parser.add_mutually_exclusive_group() - camera_type_param.add_argument('--xtion', dest='camera_type', action='store_const', const='xtion', help='Assume that the topic names are for the xtion') - camera_type_param.add_argument('--d435', dest='camera_type', action='store_const', const='d435', help='Assume that the topics are for the d435') - parser.add_argument("--file_format", default='png', help="Image file format. Default 'png'") - parser.add_argument("--image_enc", default='passthrough', - help="Image encoding. See cv_bridge tutorial for encodings. Default 'passthrough'") - parser.add_argument("--undistort", default=False, - help="Undistort input images with given camera info") - - args = parser.parse_args() - - if args.rosbag_folder is not None: - base_out_dir = args.output_dir - bag_names = [f for f in listdir(args.rosbag_folder) if isfile(join(args.rosbag_folder, f))] - for bag_name in bag_names: - if bag_name[-11:] != '_synced.bag': - continue - - args.bag_file = join(args.rosbag_folder, bag_name) - print("Extracting depth from:", bag_name) - args.output_dir = join(base_out_dir, bag_name[:-11], 'depth') - args.image_enc = 'mono16' - if args.camera_type == 'xtion': - args.image_topic = '/depth/image' - else: - args.image_topic = '/camera/aligned_depth_to_color/image_raw' - args.file_format = 'png' - main(args) - - print("Extracting color from:", bag_name) - args.output_dir = join(base_out_dir, bag_name[:-11], 'rgb') - args.image_enc = 'bgr8' - if args.camera_type == 'xtion': - args.image_topic = '/rgb/image' - else: - args.image_topic = '/camera/color/image_rect_color' - args.file_format = 'jpg' - main(args) - - - - else: - main(args) diff --git a/dataset_generation/extract_ply.py b/dataset_generation/extract_ply.py deleted file mode 100755 index 1381e7a..0000000 --- a/dataset_generation/extract_ply.py +++ /dev/null @@ -1,104 +0,0 @@ -#!/usr/bin/env python -from __future__ import print_function -import os -import argparse - -import numpy as np -import rosbag - -from ply_utils import PlyData, PlyElement -import sensor_msgs.point_cloud2 as pc2_proc - -from tqdm import tqdm -import json -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext -from pathlib import Path -import yaml - -def main(args): - - print("Extract PLY files from %s on topic %s into dir %s" % (args.bag_file, - args.pcd_topic, - args.output_dir)) - bag = rosbag.Bag(args.bag_file, "r") - count = 0 - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - field_to_type = {'x': np.float32, 'y': np.float32, 'z': np.float32, - 'r': np.uint8, 'g': np.uint8, 'b': np.uint8, - 'intensity': np.float32, - 't': np.uint32, - 'reflectivity': np.uint16, - 'ring': np.uint8, - 'noise': np.uint16, - 'range': np.uint32, - 'label': np.uint32} - - - point_fields = [] - field_names = [] - skip_nans = not args.write_nan - for topic, msg, t in bag.read_messages(topics=[args.pcd_topic]): - point_fields = msg.fields - field_names = [pf.name for pf in point_fields] - yaml_info = {} - yaml_info['data_type'] = 'pointcloud' - yaml_info['width'] = msg.width - yaml_info['height'] = msg.height - yaml_info['fields'] = field_names - yaml_info['frame_id'] = msg.header.frame_id - print("pcd_info:", yaml_info) - fname_dir = Path(args.output_dir) - with open(str(fname_dir.parent / Path(fname_dir.stem+'_info.yaml')), 'w') as f: - f.write(yaml.dump(yaml_info)) - # json.dump(yaml_info, f) - break - - for topic, msg, t in tqdm(bag.read_messages(topics=[args.pcd_topic]), - total=bag.get_message_count(args.pcd_topic)): - pts = pc2_proc.read_points(msg, field_names=field_names, skip_nans=skip_nans) - dtype = [(pf.name, field_to_type[pf.name]) for pf in point_fields] - pts_np= np.array([p for p in pts], dtype=dtype) - el = PlyElement.describe(pts_np, 'vertex') - stamp = msg.header.stamp - stamp_ns_str = str(stamp.nsecs).zfill(9) - out_file = os.path.join(args.output_dir, "frame%06i-%s_%s.ply" % (count, str(stamp.secs), - stamp_ns_str[:3])) - PlyData([el], text=args.write_ascii).write(out_file) - count += 1 - - print("Wrote {} PLY clouds to {}".format(count, args.output_dir)) - bag.close() - - return - -if __name__ == '__main__': - """Extract a folder of PLY from a rosbag. - """ - parser = argparse.ArgumentParser(description="Extract PLY from a ROS bag.", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog='''Example: - ./extract_ply.py --bag_file ~/data/rcta_box.bag --output_dir ~/data/rcta_objects/box/00001/ply/ --pcd_topic /points''') - parser.add_argument("--bag_file", default=None, help="Input ROS bag.") - parser.add_argument("--rosbag_folder", default=None, help="Folder of synchronized rosbags to extract") - parser.add_argument("--output_dir", help="Output directory.") - parser.add_argument("--pcd_topic", help="Point cloud topic.", default=None) - parser.add_argument("--write_ascii", help="Write PLY in ascii format", action="store_true") - parser.add_argument("--write_nan", help="Write NaN points", action="store_true") - - args = parser.parse_args() - - if args.rosbag_folder is not None: - base_out_dir = args.output_dir - bag_names = [f for f in listdir(args.rosbag_folder) if isfile(join(args.rosbag_folder, f))] - for bag_name in bag_names: - if bag_name[-11:] != '_synced.bag': - continue - - args.bag_file = join(args.rosbag_folder, bag_name) - args.output_dir = join(base_out_dir, bag_name[:-11], 'ply') - main(args) - else: - main(args) diff --git a/dataset_generation/label_override.yaml b/dataset_generation/label_override.yaml deleted file mode 100644 index f69f2af..0000000 --- a/dataset_generation/label_override.yaml +++ /dev/null @@ -1,3 +0,0 @@ -uphill: grass -downhill: grass -deepwater: water diff --git a/dataset_generation/merge_bags.py b/dataset_generation/merge_bags.py deleted file mode 100644 index 4c0ffee..0000000 --- a/dataset_generation/merge_bags.py +++ /dev/null @@ -1,49 +0,0 @@ -import rospy -import rosbag -import numpy as np -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import Transform, TransformStamped, Vector3, Quaternion -from std_msgs.msg import Header -import tf.transformations as tfx -from pathlib import Path -import argparse -import shutil -from tqdm import tqdm - -def main(args): - tmp_bag = "/tmp/temp.bag" - overwrite = args.out_bag in args.bags or not args.out_bag - b_out = rosbag.Bag(tmp_bag, 'w') if overwrite else rosbag.Bag(args.out_bag, 'w') - - tf_static_msg = None - tf_static_t = None - tf_static_connection_header = None - for b in args.bags: - b_in = rosbag.Bag(b) - for topic, msg, t, cnxn_hdr in tqdm(b_in.read_messages(return_connection_header=True)): - if topic == '/tf_static': - if tf_static_msg: - tf_static_msg.transforms.extend(msg.transforms) - else: - tf_static_msg = msg - tf_static_t = t - tf_static_connection_header = cnxn_hdr - else: - b_out.write(topic, msg, t) - b_in.close() - if tf_static_msg: - b_out.write('/tf_static', tf_static_msg, tf_static_t, - connection_header=tf_static_connection_header) - b_out.close() - - if overwrite: - out_bag = args.bags[-1] if not args.out_bag else args.out_bag - shutil.move(tmp_bag, out_bag) - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-b", "--bags", nargs='+', default=[], help="bags to merge") - parser.add_argument("-o", "--out_bag", default="", - help="Output ROS bag. (Default is to _overwrite_ last input bag)") - args = parser.parse_args() - main(args) diff --git a/dataset_generation/merge_jsons.py b/dataset_generation/merge_jsons.py deleted file mode 100644 index a522b7f..0000000 --- a/dataset_generation/merge_jsons.py +++ /dev/null @@ -1,260 +0,0 @@ -#!/usr/bin/env python -from __future__ import print_function -import os -import rospy -import argparse -import urllib.request -import numpy as np -import rosbag -from skimage import io -from pathlib import Path -from tqdm import tqdm -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext -import json -import csv -import cv2 -import yaml -# import tf2_ros -# from tf2_msgs.msg import TFMessage - -def hex_to_rgb(h): - h = h.lstrip('#') - hlen = len(h) - return tuple(int(h[i:i+hlen//3], 16) for i in range(0, hlen, hlen//3)) - -def hex_to_bgr(h): - return tuple(reversed(hex_to_rgb(h))) - -def get_ontology(ontology_path): - id_color_map = {} - id_rgb_map = {} - id_name_map = {} - if ontology_path.exists(): - ontology_lines = open(str(ontology_path),'rU').readlines() - for parts in csv.reader(ontology_lines[1:], quotechar='"', delimiter=',', - quoting=csv.QUOTE_ALL, skipinitialspace=True): - color_hex = parts[0] - lab_id = parts[3] - lab_name = parts[2]#.split(' ')[0] - id_color_map[int(lab_id)] = hex_to_bgr(color_hex) - id_rgb_map[int(lab_id)] = hex_to_rgb(color_hex) - id_name_map[int(lab_id)] = lab_name - - num_ontology_blocks = int(np.sqrt(len(id_color_map.keys()))) - if num_ontology_blocks**2 != len(id_color_map.keys()): - num_ontology_blocks += 1 - block_sz = 200 - ontology_img = np.zeros((block_sz*num_ontology_blocks,block_sz*num_ontology_blocks,3), dtype=np.uint8) - ont_block_idx = 0 - block_txt_map = {} - block_color_map = {} - font_type = cv2.FONT_HERSHEY_SIMPLEX - font_scale = 1 - font_thickness = 1 - for o,c in id_color_map.items(): - j_idx = ont_block_idx // num_ontology_blocks - i_idx = (ont_block_idx - j_idx * num_ontology_blocks) % num_ontology_blocks - ont_start_pt = np.array([block_sz*i_idx,block_sz*j_idx]) - ont_end_pt = ont_start_pt + np.array([block_sz,block_sz]) - ont_text_pt = (ont_start_pt + ont_end_pt - np.array([block_sz,0])) // 2 - ontology_img[ont_start_pt[1]:ont_end_pt[1],ont_start_pt[0]:ont_end_pt[0],:] = c - block_txt_map[tuple(ont_text_pt)] = id_name_map[o]+'('+str(o)+')' - block_color_map[tuple(ont_text_pt)] = c - ont_block_idx += 1 - for p,t in block_txt_map.items(): - font_sz, _ = cv2.getTextSize(t, font_type, font_scale, font_thickness) - fscale = min(font_scale, block_sz / font_sz[0]) - cv2.putText(ontology_img, t, p, font_type, fscale, (0,0,0), font_thickness) - - cv2.imwrite(str(ontology_path.with_suffix('.png')), ontology_img) - - with open(str(ontology_path.with_suffix('.yaml')), 'w') as f: - f.write(yaml.dump([id_name_map, id_rgb_map])) - - # dump_jsonl([id_name_map, id_rgb_map], str(ontology_path.with_suffix('.jsonl'))) - - return (id_color_map, id_name_map) - -def dump_jsonl(data, output_path, append=False): - """ - Write list of objects to a JSON lines file. - """ - mode = 'a+' if append else 'w' - with open(output_path, mode, encoding='utf-8') as f: - for line in data: - json_record = json.dumps(line, ensure_ascii=False) - f.write(json_record + '\n') - print('Wrote {} records to {}'.format(len(data), output_path)) - -def load_jsonl(input_path) -> list: - """ - Read list of objects from a JSON lines file. - """ - data = [] - with open(input_path, 'r', encoding='utf-8') as f: - for line in f: - data.append(json.loads(line.rstrip('\n|\r'))) - print('Loaded {} records from {}'.format(len(data), input_path)) - return data - -def stamp_str_extract(fname): - stem = Path(fname).stem - stamp_idx = stem.rfind('-') - return stem[stamp_idx+1:] - -def stamp_from_filename(fname): - stem = Path(fname).stem - msec_idx = stem.rfind('_') - msecs = float(stem[msec_idx+1:]) - sec_idx = stem[:msec_idx].rfind('-') - secs = float(stem[sec_idx+1:msec_idx]) - stamp = secs + (msecs / 1000.0) - return stamp - -def seq_in_bag(seq, bag): - bag_start = bag.get_start_time() - bag_end = bag.get_end_time() - print('bag',bag_start,bag_end,'seq',seq[0],seq[1]) - return (seq[0] >= bag_start and seq[1] <= bag_end) - -def stamp_in_seq(stamp, seq): - return (seq[0]<= stamp and seq[1] >= stamp) - -def get_associated_fname(seq, fname, args): - fname_stamp = stamp_str_extract(fname) - datadirs = [d for d in Path(seq).iterdir() if d.is_dir()] - data_files = sorted([df for dd in datadirs for df in dd.iterdir() - if df.suffix in args.data_extensions]) - for df in data_files: - if fname_stamp in str(df): - return (df.parent.stem, df.stem) - return None, None - -def main(args): - root_path = Path(args.root_dir) - collection_path = root_path / Path(args.collection_name) - - # Get sequence start/end timestamps - seq_stamp_map = {} # sequence id: (start time, end time) - stamp_seq_map = {} - sdirs = [d for d in collection_path.iterdir() if d.is_dir()] - for sd in sdirs: - if args.seq_dirs and str(sd.stem) not in args.seq_dirs: - continue - print(sd) - ddl = [d for d in Path(sd).iterdir() if d.is_dir()] - data_files = sorted([df for dd in ddl for df in dd.iterdir() - if df.suffix in args.data_extensions]) - stamps_sorted = sorted([stamp_from_filename(df) for df in data_files]) - if not stamps_sorted: - continue - timestamp_begin = stamps_sorted[0] - timestamp_end = stamps_sorted[-1] - print('ts_begin',timestamp_begin,'ts_end',timestamp_end) - seq_stamp_map[str(sd)] = (timestamp_begin, timestamp_end) - stamp_seq_map[timestamp_begin] = str(sd) - sorted_seqs = [stamp_seq_map[k] for k in sorted(stamp_seq_map.keys())] - sorted_stamps = [seq_stamp_map[sk] for sk in sorted_seqs] - print('All sequences in order: ',sorted_seqs) - print('All stamps in order', sorted_stamps) - print('Stamp differences', [sorted_stamps[i][0] - sorted_stamps[i-1][1] - for i in range(len(sorted_stamps)) if i>0]) - - # Process Appen jsonlist files - seq_json_map = {} - seq_img_map = {} - seq_ann_map = {} - for jlf in args.jsonl: - jlines = load_jsonl(jlf) - image_urls = [] - ann_urls = [] - for jl in jlines: - for jmt in jl['results']['judgments']: - if jmt['data']['broken_link'] == 'true': - continue - image_urls.append(jmt['unit_data']['image_url']) - ann_urls.append(json.loads(jmt['data']['annotation'])['url']) - timestamp_jl = stamp_from_filename(image_urls[len(image_urls) // 2]) - for sid, stamps in seq_stamp_map.items(): - if stamp_in_seq(timestamp_jl, stamps): - seq_json_map.setdefault(sid, []).append([jlf]) - seq_img_map.setdefault(sid, []).extend(image_urls) - seq_ann_map.setdefault(sid, []).extend(ann_urls) - - # Process Appen ontology file - ontology_path = collection_path / Path('ontology.csv') - id_color_map, id_name_map = get_ontology(ontology_path) - - # Extract images - for sid, fnames in seq_img_map.items(): - for i, fname in enumerate(fnames): - ext = Path(fname).suffix - ext_out = ext if ext != '.jpg' else '.png' - sensor_name, fname_stem = get_associated_fname(sid, fname, args) - if sensor_name: - out_name = collection_path / Path(sid) / Path(str(sensor_name)+'_label_id/'+fname_stem+ext_out) - if not out_name.parent.exists(): - out_name.parent.mkdir() - try: - urllib.request.urlretrieve(seq_ann_map[sid][i], str(out_name)) - if id_color_map: - color_file = collection_path / Path(sid) / Path(str(sensor_name)+'label_color/'+fname_stem+ext_out) - if not color_file.parent.exists(): - color_file.parent.mkdir() - img = io.imread(str(out_name)) - img_flat = np.reshape(img, (-1,img.shape[-1])) - img_bgr_flat = np.array([id_color_map.get(p[0],(0,0,0)) for p in img_flat], dtype=np.uint8) - img_bgr = img_bgr_flat.reshape(img.shape[0],img.shape[1],3) - cv2.imwrite(str(color_file), img_bgr) - except Exception as e: - print (e) - continue - - - # Check for bags in //bags. Associate with sequences - bag_path = collection_path / Path('bags') - bag_seq_map={} - if bag_path.exists(): - print('bag path',bag_path) - bag_names = [b for b in bag_path.iterdir() if b.suffix == '.bag' and '_synced.bag' not in str(b)] - for b in bag_names: - bag = rosbag.Bag(str(b)) - for sid, stamps in seq_stamp_map.items(): - if seq_in_bag(stamps, bag): - bag_seq_map.setdefault(b.name,[]).append([sid]) - - for b,s in bag_seq_map.items(): - seq_map_path = collection_path / Path(s) / Path('source_bag.txt') - seq_map_path.write_text(b) - - bag_seq_path = collection_path / Path('bag_sequence_map.yaml') - with open(str(bag_seq_path), 'w') as f: - f.write(yaml.dump(bag_seq_map)) - - - # For each sequence, get transform between sensors, write to file - - -if __name__ == '__main__': - """Merge jsonlines formatted files, given sequences of filenames. Filename naming convention is assumed to be - of the form -secs_msec. - """ - parser = argparse.ArgumentParser(description="Associate jsonlines files to sequences") - parser.add_argument("-r", "--root_dir", default="", - help="dataset root dir") - parser.add_argument("-n", "--collection_name", default="rellis", - help="collection name") - parser.add_argument("-s", "--seq_dirs", nargs='+',default=[], - help="sequence_dirs to process, if empty, process entire collection") - parser.add_argument("-j", "--jsonl", nargs='+', default=[], - help="Appen formatted JSONlines files.") - parser.add_argument("-e", "--data_extensions", nargs='+',default=['.png','.jpg','.ply'], - help="set of supported extensions of raw data files") - # parser.add_argument("-o", "--ontology_file", default="", - # help="Appen formatted ontology file") - # parser.add_argument("-c", "--color", action="store_true", - # help="Output label color images as well as label id images") - - args = parser.parse_args() - main(args) diff --git a/dataset_generation/merge_tf_static.py b/dataset_generation/merge_tf_static.py deleted file mode 100755 index beb35c2..0000000 --- a/dataset_generation/merge_tf_static.py +++ /dev/null @@ -1,71 +0,0 @@ -import rospy -import rosbag -import numpy as np -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import Transform, TransformStamped, Vector3, Quaternion -from std_msgs.msg import Header -import tf.transformations as tfx -from pathlib import Path -from tqdm import tqdm -import argparse -import shutil - -def main(args): - tmp_bag = "/tmp/temp.bag" - overwrite = args.out_bag in args.bags or not args.out_bag - b_out = rosbag.Bag(tmp_bag, 'w') if overwrite else rosbag.Bag(args.out_bag, 'w') - - tf_static_msg = None - tf_static_t = None - tf_static_connection_header = None - - time_min = 1e15 - bag_min = None - if not args.ref_bag: - for b in args.bags: - b_in = rosbag.Bag(b) - print 'Start time ',b_in.get_start_time() - if b_in.get_start_time() < time_min: - time_min = b_in.get_start_time() - bag_min = b - b_in.close() - args.ref_bag = bag_min - - for b in args.bags: - reference_bag = False - b_in = rosbag.Bag(b) - print b - if args.ref_bag and args.ref_bag == b: - reference_bag = True - tf_static_t = b_in.get_start_time() + 1.0 - for topic, msg, t, cnxn_hdr in tqdm(b_in.read_messages(return_connection_header=True)): - if topic == '/tf_static': - if reference_bag: - tf_static_connection_header = cnxn_hdr - if tf_static_msg: - tf_static_msg.transforms.extend(msg.transforms) - else: - tf_static_msg = msg - else: - if reference_bag: - b_out.write(topic, msg, t) - b_in.close() - - for i, tfm in enumerate(tf_static_msg.transforms): - tf_static_msg.transforms[i].header.stamp = rospy.Time.from_sec(tf_static_t) - b_out.write('/tf_static', tf_static_msg, rospy.Time.from_sec(tf_static_t), - connection_header=tf_static_connection_header) - b_out.close() - - if overwrite: - out_bag = args.bags[-1] if not args.out_bag else args.out_bag - shutil.move(tmp_bag, out_bag) - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-b", "--bags", nargs='+', default=[], help="bags to merge") - parser.add_argument("-r", "--ref_bag", default="", help="reference bag (timestamp is relative to this bag). Default is earliest time bag in the list") - parser.add_argument("-o", "--out_bag", default="", - help="Output ROS bag. (Default is to _overwrite_ last input bag)") - args = parser.parse_args() - main(args) diff --git a/dataset_generation/patch_appen.py b/dataset_generation/patch_appen.py deleted file mode 100644 index 84857f1..0000000 --- a/dataset_generation/patch_appen.py +++ /dev/null @@ -1,336 +0,0 @@ -#!/usr/bin/env python -from __future__ import print_function -import os -import rospy -import argparse -import urllib.request -import numpy as np -import rosbag -from skimage import io -from pathlib import Path -from tqdm import tqdm -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext -import json -import csv -import cv2 -import yaml -# import tf2_ros -# from tf2_msgs.msg import TFMessage - -def hex_to_rgb(h): - h = h.lstrip('#') - hlen = len(h) - return tuple(int(h[i:i+hlen//3], 16) for i in range(0, hlen, hlen//3)) - -def hex_to_bgr(h): - return tuple(reversed(hex_to_rgb(h))) - -def get_ontology(ontology_path): - id_color_map = {} - id_rgb_map = {} - id_name_map = {} - if ontology_path.exists(): - ontology_lines = open(str(ontology_path),'rU').readlines() - for parts in csv.reader(ontology_lines[1:], quotechar='"', delimiter=',', - quoting=csv.QUOTE_ALL, skipinitialspace=True): - color_hex = parts[0] - lab_id = parts[3] - lab_name = parts[2]#.split(' ')[0] - id_color_map[int(lab_id)] = hex_to_bgr(color_hex) - id_rgb_map[int(lab_id)] = hex_to_rgb(color_hex) - id_name_map[int(lab_id)] = lab_name - - num_ontology_blocks = int(np.sqrt(len(id_color_map.keys()))) - if num_ontology_blocks**2 != len(id_color_map.keys()): - num_ontology_blocks += 1 - block_sz = 200 - ontology_img = np.zeros((block_sz*num_ontology_blocks,block_sz*num_ontology_blocks,3), dtype=np.uint8) - ont_block_idx = 0 - block_txt_map = {} - block_color_map = {} - font_type = cv2.FONT_HERSHEY_SIMPLEX - font_scale = 1 - font_thickness = 1 - for o,c in id_color_map.items(): - j_idx = ont_block_idx // num_ontology_blocks - i_idx = (ont_block_idx - j_idx * num_ontology_blocks) % num_ontology_blocks - ont_start_pt = np.array([block_sz*i_idx,block_sz*j_idx]) - ont_end_pt = ont_start_pt + np.array([block_sz,block_sz]) - ont_text_pt = (ont_start_pt + ont_end_pt - np.array([block_sz,0])) // 2 - ontology_img[ont_start_pt[1]:ont_end_pt[1],ont_start_pt[0]:ont_end_pt[0],:] = c - block_txt_map[tuple(ont_text_pt)] = id_name_map[o]+'('+str(o)+')' - block_color_map[tuple(ont_text_pt)] = c - ont_block_idx += 1 - for p,t in block_txt_map.items(): - font_sz, _ = cv2.getTextSize(t, font_type, font_scale, font_thickness) - fscale = min(font_scale, block_sz / font_sz[0]) - cv2.putText(ontology_img, t, p, font_type, fscale, (0,0,0), font_thickness) - - cv2.imwrite(str(ontology_path.with_suffix('.png')), ontology_img) - - with open(str(ontology_path.with_suffix('.yaml')), 'w') as f: - f.write(yaml.dump([id_name_map, id_rgb_map])) - - # dump_jsonl([id_name_map, id_rgb_map], str(ontology_path.with_suffix('.jsonl'))) - - return (id_color_map, id_name_map) - -def dump_jsonl(data, output_path, append=False): - """ - Write list of objects to a JSON lines file. - """ - mode = 'a+' if append else 'w' - with open(output_path, mode, encoding='utf-8') as f: - for line in data: - json_record = json.dumps(line, ensure_ascii=False) - f.write(json_record + '\n') - print('Wrote {} records to {}'.format(len(data), output_path)) - -def load_jsonl(input_path) -> list: - """ - Read list of objects from a JSON lines file. - """ - data = [] - with open(input_path, 'r', encoding='utf-8') as f: - for line in f: - data.append(json.loads(line.rstrip('\n|\r'))) - print('Loaded {} records from {}'.format(len(data), input_path)) - return data - -def stamp_str_extract(fname): - stem = Path(fname).stem - stamp_idx = stem.rfind('-') - return stem[stamp_idx+1:] - -def stamp_from_filename(fname): - stem = Path(fname).stem - msec_idx = stem.rfind('_') - msecs = float(stem[msec_idx+1:]) - sec_idx = stem[:msec_idx].rfind('-') - secs = float(stem[sec_idx+1:msec_idx]) - stamp = secs + (msecs / 1000.0) - return stamp - -def old_stamp_from_filename(fname): - stem = Path(fname).stem - msec_idx = stem.rfind('_') - msecs = float(stem[msec_idx+1:]) - sec_idx = stem[:msec_idx].rfind('_') - secs = float(stem[sec_idx+1:msec_idx]) - stamp = secs + (msecs / 1000.0) - return stamp - -def seq_in_bag(seq, bag): - bag_start = bag.get_start_time() - bag_end = bag.get_end_time() - print('bag',bag_start,bag_end,'seq',seq[0],seq[1]) - return (seq[0] >= bag_start and seq[1] <= bag_end) - -def stamp_in_seq(stamp, seq): - return (seq[0]<= stamp and seq[1] >= stamp) - -def get_associated_fname(seq, fname, args): - fname_stamp = stamp_str_extract(fname) - datadirs = [d for d in Path(seq).iterdir() if d.is_dir()] - data_files = sorted([df for dd in datadirs for df in dd.iterdir() - if df.suffix in args.data_extensions]) - for df in data_files: - if fname_stamp in str(df): - return (df.parent.stem, df.stem) - return None, None - -def main(args): - root_path = Path(args.root_dir) - collection_path = root_path / Path(args.collection_name) - - # Get sequence start/end timestamps - seq_stamp_map = {} # sequence id: (start time, end time) - stamp_seq_map = {} - sdirs = [d for d in collection_path.iterdir() if d.is_dir()] - all_fnames = [] - all_stamps = [] - for sd in sdirs: - print('SD',sd) - if args.seq_dirs and str(sd.stem) not in args.seq_dirs: - continue - print(sd) - ddl = [d for d in Path(sd).iterdir() if d.is_dir() and 'label_id' in str(d.stem)] - print ('DDL',ddl) - data_files = sorted([df for dd in ddl for df in dd.iterdir() - if df.suffix in args.data_extensions]) - stamps_sorted = sorted([stamp_from_filename(df) for df in data_files]) - fnames = [df for df in data_files] - stamps = [stamp_from_filename(df) for df in fnames] - if not stamps_sorted: - continue - all_fnames.extend([str(f) for f in fnames]) - all_stamps.extend(stamps) - timestamp_begin = stamps_sorted[0] - timestamp_end = stamps_sorted[-1] - print('ts_begin',timestamp_begin,'ts_end',timestamp_end) - seq_stamp_map[str(sd)] = (timestamp_begin, timestamp_end) - stamp_seq_map[timestamp_begin] = str(sd) - sorted_seqs = [stamp_seq_map[k] for k in sorted(stamp_seq_map.keys())] - sorted_stamps = [seq_stamp_map[sk] for sk in sorted_seqs] - print('All sequences in order: ',sorted_seqs) - print('All stamps in order', sorted_stamps) - print('Stamp differences', [sorted_stamps[i][0] - sorted_stamps[i-1][1] - for i in range(len(sorted_stamps)) if i>0]) - print('all',len(all_fnames)) - - sort_idxs = np.argsort(all_stamps) - all_stamps = np.array(all_stamps)[sort_idxs].tolist() - all_fnames = np.array(all_fnames)[sort_idxs].tolist() - print('ALL FNAMES',all_fnames[:10000]) - - - # The worst - old_fnames = [] - old_stamps = [] - tmppath = Path("/data/tmp/mod3") - old_fnames = [fn for fn in tmppath.glob('**/*.png')] - old_stamps = [old_stamp_from_filename(df) for df in old_fnames] - sorted_old = np.argsort(old_stamps) - old_stamps = np.array(old_stamps)[sorted_old].tolist() - old_fnames = np.array(old_fnames)[sorted_old].tolist() - - # Process Appen jsonlist files - seq_json_map = {} - seq_img_map = {} - seq_ann_map = {} - img_ann_map = {} - - for jlf in args.jsonl: - jlines = load_jsonl(jlf) - image_urls = [] - ann_urls = [] - for jl in tqdm(jlines): - for jmt in jl['results']['judgments']: - if jmt['data']['broken_link'] == 'true': - continue - image_urls.append(jmt['unit_data']['image_url']) - stamp = stamp_from_filename(jmt['unit_data']['image_url']) - ann_urls.append(json.loads(jmt['data']['annotation'])['url']) - img_ann_map[stamp] = json.loads(jmt['data']['annotation'])['url'] - timestamp_jl = stamp_from_filename(image_urls[len(image_urls) // 2]) - for sid, stamps in seq_stamp_map.items(): - if stamp_in_seq(timestamp_jl, stamps): - seq_json_map.setdefault(sid, []).append([jlf]) - seq_img_map.setdefault(sid, []).extend(image_urls) - seq_ann_map.setdefault(sid, []).extend(ann_urls) - assert(len(image_urls) == len(ann_urls)) - - missing=[] - urls = [] - for k,v in img_ann_map.items(): - if k not in all_stamps and k not in old_stamps: - for i,s in enumerate(all_stamps): - if s > k: - if i == 0: - break - missing.append(k) - fmti = str(i).zfill(6) - ns_str = str(round(1000*(k - int(k)))) - out_stem = 'frame'+fmti+'_'+str(int(k))+'_'+ns_str+'.png' - out_name = tmppath / Path(out_stem) - print('URL REQ',v + ' ' + out_stem) - # urllib.request.urlretrieve(v, str(out_name)) - urls.append(v + ' ' + out_stem) - break - - with open(str(tmppath / 'urls.txt'), 'w') as f: - f.writelines("%s\n" % u for u in urls) - with open(str(tmppath / 'ann_map.yaml'), 'w') as f: - f.write(yaml.dump(img_ann_map)) - print('len missing',len(missing)) - print('len all',len(img_ann_map.keys())) - assert(False) - - # Process Appen ontology file - ontology_path = collection_path / Path('ontology.csv') - id_color_map, id_name_map = get_ontology(ontology_path) - - # # Extract images - # for sid, fnames in seq_img_map.items(): - # print('Seq:',sid) - # for i, fname in tqdm(enumerate(fnames)): - # ext = Path(fname).suffix - # ext_out = ext if ext != '.jpg' else '.png' - # sensor_name, fname_stem = get_associated_fname(sid, fname, args) - # if sensor_name: - # out_name = collection_path / Path(sid) / Path(str(sensor_name)+'_label_id/'+fname_stem+ext_out) - # if not out_name.parent.exists(): - # out_name.parent.mkdir() - # try: - # urllib.request.urlretrieve(seq_ann_map[sid][i], str(out_name)) - # if id_color_map: - # color_file = collection_path / Path(sid) / Path(str(sensor_name)+'_label_color/'+fname_stem+ext_out) - # if not color_file.parent.exists(): - # color_file.parent.mkdir() - # img = io.imread(str(out_name)) - # img_flat = np.reshape(img, (-1,img.shape[-1])) - # img_bgr_flat = np.array([id_color_map.get(p[0],(0,0,0)) for p in img_flat], dtype=np.uint8) - # img_bgr = img_bgr_flat.reshape(img.shape[0],img.shape[1],3) - # cv2.imwrite(str(color_file), img_bgr) - # except Exception as e: - # try: - # urllib.request.urlretrieve(seq_ann_map[sid][i], str(out_name)) - # if id_color_map: - # color_file = collection_path / Path(sid) / Path(str(sensor_name)+'_label_color/'+fname_stem+ext_out) - # if not color_file.parent.exists(): - # color_file.parent.mkdir() - # img = io.imread(str(out_name)) - # img_flat = np.reshape(img, (-1,img.shape[-1])) - # img_bgr_flat = np.array([id_color_map.get(p[0],(0,0,0)) for p in img_flat], dtype=np.uint8) - # img_bgr = img_bgr_flat.reshape(img.shape[0],img.shape[1],3) - # cv2.imwrite(str(color_file), img_bgr) - # except Exception as e2: - # print (e2) - # continue - - - # Check for bags in //bags. Associate with sequences - bag_path = collection_path / Path('bags') - bag_seq_map={} - if bag_path.exists(): - print('bag path',bag_path) - bag_names = [b for b in bag_path.iterdir() if b.suffix == '.bag' and '_synced.bag' not in str(b)] - for b in bag_names: - bag = rosbag.Bag(str(b)) - for sid, stamps in seq_stamp_map.items(): - if seq_in_bag(stamps, bag): - bag_seq_map[b.name] = sid # assumes 1-1 bag/seq mapping (any split bags have already been merged) - # bag_seq_map.setdefault(b.name,[]).append([sid]) - - bag_seq_path = collection_path / Path('bag_sequence_map.yaml') - bag_seq_stem_map = {k:Path(v).stem for k,v in bag_seq_map.items()} - with open(str(bag_seq_path), 'w') as f: - f.write(yaml.dump(bag_seq_stem_map)) - - for b,s in bag_seq_map.items(): - seq_map_path = collection_path / Path(s) / Path('source_bag.txt') - seq_map_path.write_text(b) - - # For each sequence, get transform between sensors, write to file - - -if __name__ == '__main__': - parser = argparse.ArgumentParser(description="Associate jsonlines files to sequences") - parser.add_argument("-r", "--root_dir", default="", - help="dataset root dir") - parser.add_argument("-n", "--collection_name", default="rellis", - help="collection name") - parser.add_argument("-s", "--seq_dirs", nargs='+',default=[], - help="sequence_dirs to process, if empty, process entire collection") - parser.add_argument("-j", "--jsonl", nargs='+', default=[], - help="Appen formatted JSONlines files.") - parser.add_argument("-e", "--data_extensions", nargs='+',default=['.png','.jpg','.ply'], - help="set of supported extensions of raw data files") - # parser.add_argument("-o", "--ontology_file", default="", - # help="Appen formatted ontology file") - # parser.add_argument("-c", "--color", action="store_true", - # help="Output label color images as well as label id images") - - args = parser.parse_args() - main(args) diff --git a/dataset_generation/patch_bag_caminfo.py b/dataset_generation/patch_bag_caminfo.py deleted file mode 100755 index 338541c..0000000 --- a/dataset_generation/patch_bag_caminfo.py +++ /dev/null @@ -1,73 +0,0 @@ -import rospy -import rosbag -import numpy as np -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import Transform, TransformStamped, Vector3, Quaternion -from std_msgs.msg import Header -from sensor_msgs.msg import CameraInfo -import tf.transformations as tfx -from pathlib import Path -import argparse -import shutil -from tqdm import tqdm - -def main(args): - tmp_bag = "/tmp/temp.bag" - overwrite = args.out_bag == args.bag or not args.out_bag - b_out = (rosbag.Bag(tmp_bag, 'w') if overwrite else rosbag.Bag(args.out_bag, 'w')) - b_in = rosbag.Bag(args.bag) - b_cam_info = rosbag.Bag(args.cam_info_bag) - - cam_ns_msg_map = {} - for topic, msg, t in b_cam_info.read_messages(): - if 'CameraInfo' in str(type(msg)): - if "nerian" in topic and "right" in topic: - topic = "/nerian/right/camera_info" - if "nerian" in topic and "left" in topic: - topic = "/nerian/left/camera_info" - - topic_path = Path(topic) - cam_ns_msg_map[str(topic_path.parent)] = msg - print (topic) - b_cam_info.close() - - # print 'cam_ns_map',cam_ns_msg_map - topics = b_in.get_type_and_topic_info()[1].keys() - for topic, msg, t, cnxn_hdr in tqdm(b_in.read_messages(return_connection_header=True)): - if 'Image' in str(type(msg)): - if "nerian" in topic and "right" in topic: - topic = "/nerian/right/image_raw" - if "nerian" in topic and "left" in topic: - topic = "/nerian/left/image_raw" - - topic_path = Path(topic) - if (str(topic_path.parent) in cam_ns_msg_map and - str(topic_path.parent)+"/camera_info" not in topics): - cam_ns_msg_map[str(topic_path.parent)].header.stamp = msg.header.stamp - b_out.write(str(topic_path.parent)+"/camera_info", - cam_ns_msg_map[str(topic_path.parent)], t) - if 'CameraInfo' in str(type(msg)): - topic_path = Path(topic) - if (str(topic_path.parent) in cam_ns_msg_map): - cam_ns_msg_map[str(topic_path.parent)].header.stamp = msg.header.stamp - b_out.write(topic, cam_ns_msg_map[str(topic_path.parent)], t) - else: - b_out.write(topic, msg, t) - elif 'tf_static' in topic: - b_out.write(topic, msg, t, connection_header=cnxn_hdr) - else: - b_out.write(topic, msg, t) - b_in.close() - b_out.close() - - if overwrite: - shutil.move(tmp_bag, args.bag) - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-b", "--bag", default="") - parser.add_argument('-c', "--cam_info_bag", default="") - parser.add_argument("-o", "--out_bag", default="", - help="Output ROS bag. (Default is to _overwrite_ input bag)") - args = parser.parse_args() - main(args) diff --git a/dataset_generation/patch_bag_tf.py b/dataset_generation/patch_bag_tf.py deleted file mode 100755 index e5e7b05..0000000 --- a/dataset_generation/patch_bag_tf.py +++ /dev/null @@ -1,67 +0,0 @@ -import rospy -import rosbag -import numpy as np -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import Transform, TransformStamped, Vector3, Quaternion -from std_msgs.msg import Header -import tf.transformations as tfx -from pathlib import Path -import argparse -import shutil - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-b", "--bag", default="") - parser.add_argument("-o", "--out_bag", default="", - help="Output ROS bag. (Default is to _overwrite_ input bag)") - parser.add_argument('-p', "--parent", default="ouster1/os1_lidar") - parser.add_argument('-c', "--child", default="pylon_camera") - args = parser.parse_args() - - b_path = Path(args.bag) - b_in = rosbag.Bag(str(b_path)) - tmp_bag = '/tmp/tmp.bag' - b_out=rosbag.Bag(tmp_bag, 'w') - # b_out = rosbag.Bag(str(Path(args.out_dir) / Path(b_path.stem))+'.bag', 'w') - - updated_tf = False - has_static = b_in.get_message_count(topic_filters=['/tf_static']) > 0 - R = np.array([[0.0299636, 0.999547, -0.00265867], - [-0.00909438, -0.00238713, -0.999956], - [-0.99951, 0.0299864, 0.00901874]]) - t = np.array([-0.0537767, -0.269452, -0.322621]) - T = tfx.identity_matrix() - T[:3,:3] = R - T[:3,3] = t - t = tfx.translation_from_matrix(tfx.inverse_matrix(T)) - q = tfx.quaternion_from_matrix(tfx.inverse_matrix(T)) - transform_msg = Transform() - transform_msg.translation = Vector3(x=t[0], y=t[1], z=t[2]) - transform_msg.rotation = Quaternion(x=q[0], y=q[1], z=q[2], w=q[3]) - - for topic, msg, t, cnxn_hdr in b_in.read_messages(return_connection_header=True): - if not has_static and not updated_tf: - tf_msg = TFMessage() - tf_stamped = TransformStamped() - tf_stamped.header.frame_id = args.parent - tf_stamped.header.stamp = t + rospy.Duration(0.5) - tf_stamped.child_frame_id = args.child - tf_stamped.transform = transform_msg - tf_msg.transforms = [tf_stamped] - b_out.write('/tf_static', tf_msg, t + rospy.Duration(0.5)) - updated_tf = True - if has_static and topic == '/tf_static': - tf_msg = msg - tf_stamped = TransformStamped() - tf_stamped.header.frame_id = args.parent - tf_stamped.header.stamp = t + rospy.Duration(0.5) - tf_stamped.child_frame_id = args.child - tf_stamped.transform = transform_msg - tf_msg.transforms.append(tf_stamped) - msg = tf_msg - b_out.write(topic, msg, t, connection_header=cnxn_hdr) - b_in.close() - b_out.close() - - out_bag = args.out_bag if args.out_bag else args.bag - shutil.move(tmp_bag, out_bag) diff --git a/dataset_generation/ply_utils.py b/dataset_generation/ply_utils.py deleted file mode 100644 index d822584..0000000 --- a/dataset_generation/ply_utils.py +++ /dev/null @@ -1,999 +0,0 @@ -# Copyright 2014 Darsh Ranjan -# -# This file is part of python-plyfile. -# -# python-plyfile is free software: you can redistribute it and/or -# modify it under the terms of the GNU General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# python-plyfile is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with python-plyfile. If not, see -# . - -from itertools import islice as _islice - -import numpy as _np -from sys import byteorder as _byteorder - - -try: - _range = xrange -except NameError: - _range = range - - -# Many-many relation -_data_type_relation = [ - ('int8', 'i1'), - ('char', 'i1'), - ('uint8', 'u1'), - ('uchar', 'b1'), - ('uchar', 'u1'), - ('int16', 'i2'), - ('short', 'i2'), - ('uint16', 'u2'), - ('ushort', 'u2'), - ('int32', 'i4'), - ('int', 'i4'), - ('uint32', 'u4'), - ('uint', 'u4'), - ('float32', 'f4'), - ('float', 'f4'), - ('float64', 'f8'), - ('double', 'f8') -] - -_data_types = dict(_data_type_relation) -_data_type_reverse = dict((b, a) for (a, b) in _data_type_relation) - -_types_list = [] -_types_set = set() -for (_a, _b) in _data_type_relation: - if _a not in _types_set: - _types_list.append(_a) - _types_set.add(_a) - if _b not in _types_set: - _types_list.append(_b) - _types_set.add(_b) - - -_byte_order_map = { - 'ascii': '=', - 'binary_little_endian': '<', - 'binary_big_endian': '>' -} - -_byte_order_reverse = { - '<': 'binary_little_endian', - '>': 'binary_big_endian' -} - -_native_byte_order = {'little': '<', 'big': '>'}[_byteorder] - - -def _lookup_type(type_str): - if type_str not in _data_type_reverse: - try: - type_str = _data_types[type_str] - except KeyError: - raise ValueError("field type %r not in %r" % - (type_str, _types_list)) - - return _data_type_reverse[type_str] - - -def make2d(array, cols=None, dtype=None): - ''' - Make a 2D array from an array of arrays. The `cols' and `dtype' - arguments can be omitted if the array is not empty. - - ''' - if not len(array): - if cols is None or dtype is None: - raise RuntimeError( - "cols and dtype must be specified for empty array" - ) - return _np.empty((0, cols), dtype=dtype) - return _np.vstack(array) - - -class _PlyHeaderParser(object): - def __init__(self): - self.format = None - self.elements = [] - self.comments = [] - self.obj_info = [] - self.lines = 0 - self._allowed = ['ply'] - - def consume(self, raw_line): - self.lines += 1 - if not raw_line: - self._error("early end-of-file") - - line = raw_line.decode('ascii').strip() - try: - keyword = line.split(None, 1)[0] - except IndexError: - self._error() - - if keyword not in self._allowed: - self._error("expected one of {%s}" % - ", ".join(self._allowed)) - - getattr(self, 'parse_' + keyword)(line[len(keyword)+1:]) - return self._allowed - - def _error(self, message="parse error"): - raise PlyHeaderParseError(message, self.lines) - - def parse_ply(self, data): - if data: - self._error("unexpected characters after 'ply'") - self._allowed = ['format', 'comment', 'obj_info'] - - def parse_format(self, data): - fields = data.strip().split() - if len(fields) != 2: - self._error("expected \"format {format} 1.0\"") - - self.format = fields[0] - if self.format not in _byte_order_map: - self._error("don't understand format %r" % format) - - if fields[1] != '1.0': - self._error("expected version '1.0'") - - self._allowed = ['element', 'comment', 'obj_info', 'end_header'] - - def parse_comment(self, data): - if not self.elements: - self.comments.append(data) - else: - self.elements[-1][3].append(data) - - def parse_obj_info(self, data): - self.obj_info.append(data) - - def parse_element(self, data): - fields = data.strip().split() - if len(fields) != 2: - self._error("expected \"element {name} {count}\"") - - name = fields[0] - try: - count = int(fields[1]) - except ValueError: - self._error("expected integer count") - - self.elements.append((name, [], count, [])) - self._allowed = ['element', 'comment', 'property', 'end_header'] - - def parse_property(self, data): - properties = self.elements[-1][1] - fields = data.strip().split() - if len(fields) < 2: - self._error("bad 'property' line") - - if fields[0] == 'list': - if len(fields) != 4: - self._error("expected \"property list " - "{len_type} {val_type} {name}\"") - - try: - properties.append( - PlyListProperty(fields[3], fields[1], fields[2]) - ) - except ValueError as e: - self._error(str(e)) - - else: - if len(fields) != 2: - self._error("expected \"property {type} {name}\"") - - try: - properties.append( - PlyProperty(fields[1], fields[0]) - ) - except ValueError as e: - self._error(str(e)) - - def parse_end_header(self, data): - if data: - self._error("unexpected data after 'end_header'") - self._allowed = [] - - -class PlyParseError(Exception): - - ''' - Base class for PLY parsing errors. - - ''' - - pass - - -class PlyElementParseError(PlyParseError): - - ''' - Raised when a PLY element cannot be parsed. - - The attributes `element', `row', `property', and `message' give - additional information. - - ''' - - def __init__(self, message, element=None, row=None, prop=None): - self.message = message - self.element = element - self.row = row - self.prop = prop - - s = '' - if self.element: - s += 'element %r: ' % self.element.name - if self.row is not None: - s += 'row %d: ' % self.row - if self.prop: - s += 'property %r: ' % self.prop.name - s += self.message - - Exception.__init__(self, s) - - def __repr__(self): - return ('%s(%r, element=%r, row=%r, prop=%r)' % - (self.__class__.__name__, - self.message, self.element, self.row, self.prop)) - - -class PlyHeaderParseError(PlyParseError): - - ''' - Raised when a PLY header cannot be parsed. - - The attribute `line' provides additional information. - - ''' - - def __init__(self, message, line=None): - self.message = message - self.line = line - - s = '' - if self.line: - s += 'line %r: ' % self.line - s += self.message - - Exception.__init__(self, s) - - def __repr__(self): - return ('%s(%r, line=%r)' % - (self.__class__.__name__, - self.message, self.line)) - - -class PlyData(object): - - ''' - PLY file header and data. - - A PlyData instance is created in one of two ways: by the static - method PlyData.read (to read a PLY file), or directly from __init__ - given a sequence of elements (which can then be written to a PLY - file). - - ''' - - def __init__(self, elements=[], text=False, byte_order='=', - comments=[], obj_info=[]): - ''' - elements: sequence of PlyElement instances. - - text: whether the resulting PLY file will be text (True) or - binary (False). - - byte_order: '<' for little-endian, '>' for big-endian, or '=' - for native. This is only relevant if `text' is False. - - comments: sequence of strings that will be placed in the header - between the 'ply' and 'format ...' lines. - - obj_info: like comments, but will be placed in the header with - "obj_info ..." instead of "comment ...". - - ''' - if byte_order == '=' and not text: - byte_order = _native_byte_order - - self.byte_order = byte_order - self.text = text - - self.comments = comments - self.obj_info = obj_info - self.elements = elements - - def _get_elements(self): - return self._elements - - def _set_elements(self, elements): - self._elements = tuple(elements) - self._index() - - elements = property(_get_elements, _set_elements) - - def _get_byte_order(self): - return self._byte_order - - def _set_byte_order(self, byte_order): - if byte_order not in ['<', '>', '=']: - raise ValueError("byte order must be '<', '>', or '='") - - self._byte_order = byte_order - - byte_order = property(_get_byte_order, _set_byte_order) - - def _index(self): - self._element_lookup = dict((elt.name, elt) for elt in - self._elements) - if len(self._element_lookup) != len(self._elements): - raise ValueError("two elements with same name") - - def _get_comments(self): - return list(self._comments) - - def _set_comments(self, comments): - _check_comments(comments) - self._comments = list(comments) - - comments = property(_get_comments, _set_comments) - - def _get_obj_info(self): - return list(self._obj_info) - - def _set_obj_info(self, obj_info): - _check_comments(obj_info) - self._obj_info = list(obj_info) - - obj_info = property(_get_obj_info, _set_obj_info) - - @staticmethod - def _parse_header(stream): - ''' - Parse a PLY header from a readable file-like stream. - - ''' - parser = _PlyHeaderParser() - while parser.consume(stream.readline()): - pass - - return PlyData( - [PlyElement(*e) for e in parser.elements], - parser.format == 'ascii', - _byte_order_map[parser.format], - parser.comments, - parser.obj_info - ) - - @staticmethod - def read(stream): - ''' - Read PLY data from a readable file-like object or filename. - - ''' - (must_close, stream) = _open_stream(stream, 'read') - try: - data = PlyData._parse_header(stream) - for elt in data: - elt._read(stream, data.text, data.byte_order) - finally: - if must_close: - stream.close() - - return data - - def write(self, stream): - ''' - Write PLY data to a writeable file-like object or filename. - - ''' - (must_close, stream) = _open_stream(stream, 'write') - try: - stream.write(self.header.encode('ascii')) - stream.write(b'\n') - for elt in self: - elt._write(stream, self.text, self.byte_order) - finally: - if must_close: - stream.close() - - @property - def header(self): - ''' - Provide PLY-formatted metadata for the instance. - - ''' - lines = ['ply'] - - if self.text: - lines.append('format ascii 1.0') - else: - lines.append('format ' + - _byte_order_reverse[self.byte_order] + - ' 1.0') - - # Some information is lost here, since all comments are placed - # between the 'format' line and the first element. - for c in self.comments: - lines.append('comment ' + c) - - for c in self.obj_info: - lines.append('obj_info ' + c) - - lines.extend(elt.header for elt in self.elements) - lines.append('end_header') - return '\n'.join(lines) - - def __iter__(self): - return iter(self.elements) - - def __len__(self): - return len(self.elements) - - def __contains__(self, name): - return name in self._element_lookup - - def __getitem__(self, name): - return self._element_lookup[name] - - def __str__(self): - return self.header - - def __repr__(self): - return ('PlyData(%r, text=%r, byte_order=%r, ' - 'comments=%r, obj_info=%r)' % - (self.elements, self.text, self.byte_order, - self.comments, self.obj_info)) - - -def _open_stream(stream, read_or_write): - if hasattr(stream, read_or_write): - return (False, stream) - try: - return (True, open(stream, read_or_write[0] + 'b')) - except TypeError: - raise RuntimeError("expected open file or filename") - - -class PlyElement(object): - - ''' - PLY file element. - - A client of this library doesn't normally need to instantiate this - directly, so the following is only for the sake of documenting the - internals. - - Creating a PlyElement instance is generally done in one of two ways: - as a byproduct of PlyData.read (when reading a PLY file) and by - PlyElement.describe (before writing a PLY file). - - ''' - - def __init__(self, name, properties, count, comments=[]): - ''' - This is not part of the public interface. The preferred methods - of obtaining PlyElement instances are PlyData.read (to read from - a file) and PlyElement.describe (to construct from a numpy - array). - - ''' - _check_name(name) - self._name = str(name) - self._count = count - - self._properties = tuple(properties) - self._index() - - self.comments = comments - - self._have_list = any(isinstance(p, PlyListProperty) - for p in self.properties) - - @property - def count(self): - return self._count - - def _get_data(self): - return self._data - - def _set_data(self, data): - self._data = data - self._count = len(data) - self._check_sanity() - - data = property(_get_data, _set_data) - - def _check_sanity(self): - for prop in self.properties: - if prop.name not in self._data.dtype.fields: - raise ValueError("dangling property %r" % prop.name) - - def _get_properties(self): - return self._properties - - def _set_properties(self, properties): - self._properties = tuple(properties) - self._check_sanity() - self._index() - - properties = property(_get_properties, _set_properties) - - def _get_comments(self): - return list(self._comments) - - def _set_comments(self, comments): - _check_comments(comments) - self._comments = list(comments) - - comments = property(_get_comments, _set_comments) - - def _index(self): - self._property_lookup = dict((prop.name, prop) - for prop in self._properties) - if len(self._property_lookup) != len(self._properties): - raise ValueError("two properties with same name") - - def ply_property(self, name): - return self._property_lookup[name] - - @property - def name(self): - return self._name - - def dtype(self, byte_order='='): - ''' - Return the numpy dtype of the in-memory representation of the - data. (If there are no list properties, and the PLY format is - binary, then this also accurately describes the on-disk - representation of the element.) - - ''' - return _np.dtype([(prop.name, prop.dtype(byte_order)) - for prop in self.properties]) - - @staticmethod - def describe(data, name, len_types={}, val_types={}, - comments=[]): - ''' - Construct a PlyElement from an array's metadata. - - len_types and val_types can be given as mappings from list - property names to type strings (like 'u1', 'f4', etc., or - 'int8', 'float32', etc.). These can be used to define the length - and value types of list properties. List property lengths - always default to type 'u1' (8-bit unsigned integer), and value - types default to 'i4' (32-bit integer). - - ''' - if not isinstance(data, _np.ndarray): - raise TypeError("only numpy arrays are supported") - - if len(data.shape) != 1: - raise ValueError("only one-dimensional arrays are " - "supported") - - count = len(data) - - properties = [] - descr = data.dtype.descr - - for t in descr: - if not isinstance(t[1], str): - raise ValueError("nested records not supported") - - if not t[0]: - raise ValueError("field with empty name") - - if len(t) != 2 or t[1][1] == 'O': - # non-scalar field, which corresponds to a list - # property in PLY. - - if t[1][1] == 'O': - if len(t) != 2: - raise ValueError("non-scalar object fields not " - "supported") - - len_str = _data_type_reverse[len_types.get(t[0], 'u1')] - if t[1][1] == 'O': - val_type = val_types.get(t[0], 'i4') - val_str = _lookup_type(val_type) - else: - val_str = _lookup_type(t[1][1:]) - - prop = PlyListProperty(t[0], len_str, val_str) - else: - val_str = _lookup_type(t[1][1:]) - prop = PlyProperty(t[0], val_str) - - properties.append(prop) - - elt = PlyElement(name, properties, count, comments) - elt.data = data - - return elt - - def _read(self, stream, text, byte_order): - ''' - Read the actual data from a PLY file. - - ''' - dtype = self.dtype(byte_order) - if text: - self._read_txt(stream) - elif _can_mmap(stream) and not self._have_list: - # Loading the data is straightforward. We will memory map - # the file in copy-on-write mode. - num_bytes = self.count * dtype.itemsize - offset = stream.tell() - stream.seek(0, 2) - max_bytes = stream.tell() - offset - if max_bytes < num_bytes: - raise PlyElementParseError("early end-of-file", self, - max_bytes // dtype.itemsize) - self._data = _np.memmap(stream, dtype, - 'c', offset, self.count) - # Fix stream position - stream.seek(offset + self.count * dtype.itemsize) - else: - # A simple load is impossible. - self._read_bin(stream, byte_order) - - self._check_sanity() - - def _write(self, stream, text, byte_order): - ''' - Write the data to a PLY file. - - ''' - if text: - self._write_txt(stream) - else: - if self._have_list: - # There are list properties, so serialization is - # slightly complicated. - self._write_bin(stream, byte_order) - else: - # no list properties, so serialization is - # straightforward. - stream.write(self.data.astype(self.dtype(byte_order), - copy=False).data) - - def _read_txt(self, stream): - ''' - Load a PLY element from an ASCII-format PLY file. The element - may contain list properties. - - ''' - self._data = _np.empty(self.count, dtype=self.dtype()) - - k = 0 - for line in _islice(iter(stream.readline, b''), self.count): - fields = iter(line.strip().split()) - for prop in self.properties: - try: - self._data[prop.name][k] = prop._from_fields(fields) - except StopIteration: - raise PlyElementParseError("early end-of-line", - self, k, prop) - except ValueError: - raise PlyElementParseError("malformed input", - self, k, prop) - try: - next(fields) - except StopIteration: - pass - else: - raise PlyElementParseError("expected end-of-line", - self, k) - k += 1 - - if k < self.count: - del self._data - raise PlyElementParseError("early end-of-file", self, k) - - def _write_txt(self, stream): - ''' - Save a PLY element to an ASCII-format PLY file. The element may - contain list properties. - - ''' - for rec in self.data: - fields = [] - for prop in self.properties: - fields.extend(prop._to_fields(rec[prop.name])) - - _np.savetxt(stream, [fields], '%.18g', newline='\n') - - def _read_bin(self, stream, byte_order): - ''' - Load a PLY element from a binary PLY file. The element may - contain list properties. - - ''' - self._data = _np.empty(self.count, dtype=self.dtype(byte_order)) - - for k in _range(self.count): - for prop in self.properties: - try: - self._data[prop.name][k] = \ - prop._read_bin(stream, byte_order) - except StopIteration: - raise PlyElementParseError("early end-of-file", - self, k, prop) - - def _write_bin(self, stream, byte_order): - ''' - Save a PLY element to a binary PLY file. The element may - contain list properties. - - ''' - for rec in self.data: - for prop in self.properties: - prop._write_bin(rec[prop.name], stream, byte_order) - - @property - def header(self): - ''' - Format this element's metadata as it would appear in a PLY - header. - - ''' - lines = ['element %s %d' % (self.name, self.count)] - - # Some information is lost here, since all comments are placed - # between the 'element' line and the first property definition. - for c in self.comments: - lines.append('comment ' + c) - - lines.extend(list(map(str, self.properties))) - - return '\n'.join(lines) - - def __getitem__(self, key): - return self.data[key] - - def __setitem__(self, key, value): - self.data[key] = value - - def __str__(self): - return self.header - - def __repr__(self): - return ('PlyElement(%r, %r, count=%d, comments=%r)' % - (self.name, self.properties, self.count, - self.comments)) - - -def _check_comments(comments): - for comment in comments: - for char in comment: - if not 0 <= ord(char) < 128: - raise ValueError("non-ASCII character in comment") - if char == '\n': - raise ValueError("embedded newline in comment") - - -class PlyProperty(object): - - ''' - PLY property description. This class is pure metadata; the data - itself is contained in PlyElement instances. - - ''' - - def __init__(self, name, val_dtype): - _check_name(name) - self._name = str(name) - self.val_dtype = val_dtype - - def _get_val_dtype(self): - return self._val_dtype - - def _set_val_dtype(self, val_dtype): - self._val_dtype = _data_types[_lookup_type(val_dtype)] - - val_dtype = property(_get_val_dtype, _set_val_dtype) - - @property - def name(self): - return self._name - - def dtype(self, byte_order='='): - ''' - Return the numpy dtype description for this property (as a tuple - of strings). - - ''' - return byte_order + self.val_dtype - - def _from_fields(self, fields): - ''' - Parse from generator. Raise StopIteration if the property could - not be read. - - ''' - return _np.dtype(self.dtype()).type(next(fields)) - - def _to_fields(self, data): - ''' - Return generator over one item. - - ''' - yield _np.dtype(self.dtype()).type(data) - - def _read_bin(self, stream, byte_order): - ''' - Read data from a binary stream. Raise StopIteration if the - property could not be read. - - ''' - try: - return _read_array(stream, self.dtype(byte_order), 1)[0] - except IndexError: - raise StopIteration - - def _write_bin(self, data, stream, byte_order): - ''' - Write data to a binary stream. - - ''' - _write_array(stream, _np.dtype(self.dtype(byte_order)).type(data)) - - def __str__(self): - val_str = _data_type_reverse[self.val_dtype] - return 'property %s %s' % (val_str, self.name) - - def __repr__(self): - return 'PlyProperty(%r, %r)' % (self.name, - _lookup_type(self.val_dtype)) - - -class PlyListProperty(PlyProperty): - - ''' - PLY list property description. - - ''' - - def __init__(self, name, len_dtype, val_dtype): - PlyProperty.__init__(self, name, val_dtype) - - self.len_dtype = len_dtype - - def _get_len_dtype(self): - return self._len_dtype - - def _set_len_dtype(self, len_dtype): - self._len_dtype = _data_types[_lookup_type(len_dtype)] - - len_dtype = property(_get_len_dtype, _set_len_dtype) - - def dtype(self, byte_order='='): - ''' - List properties always have a numpy dtype of "object". - - ''' - return '|O' - - def list_dtype(self, byte_order='='): - ''' - Return the pair (len_dtype, val_dtype) (both numpy-friendly - strings). - - ''' - return (byte_order + self.len_dtype, - byte_order + self.val_dtype) - - def _from_fields(self, fields): - (len_t, val_t) = self.list_dtype() - - n = int(_np.dtype(len_t).type(next(fields))) - - data = _np.loadtxt(list(_islice(fields, n)), val_t, ndmin=1) - if len(data) < n: - raise StopIteration - - return data - - def _to_fields(self, data): - ''' - Return generator over the (numerical) PLY representation of the - list data (length followed by actual data). - - ''' - (len_t, val_t) = self.list_dtype() - - data = _np.asarray(data, dtype=val_t).ravel() - - yield _np.dtype(len_t).type(data.size) - for x in data: - yield x - - def _read_bin(self, stream, byte_order): - (len_t, val_t) = self.list_dtype(byte_order) - - try: - n = _read_array(stream, _np.dtype(len_t), 1)[0] - except IndexError: - raise StopIteration - - data = _read_array(stream, _np.dtype(val_t), n) - if len(data) < n: - raise StopIteration - - return data - - def _write_bin(self, data, stream, byte_order): - ''' - Write data to a binary stream. - - ''' - (len_t, val_t) = self.list_dtype(byte_order) - - data = _np.asarray(data, dtype=val_t).ravel() - - _write_array(stream, _np.array(data.size, dtype=len_t)) - _write_array(stream, data) - - def __str__(self): - len_str = _data_type_reverse[self.len_dtype] - val_str = _data_type_reverse[self.val_dtype] - return 'property list %s %s %s' % (len_str, val_str, self.name) - - def __repr__(self): - return ('PlyListProperty(%r, %r, %r)' % - (self.name, - _lookup_type(self.len_dtype), - _lookup_type(self.val_dtype))) - - -def _check_name(name): - for char in name: - if not 0 <= ord(char) < 128: - raise ValueError("non-ASCII character in name %r" % name) - if char.isspace(): - raise ValueError("space character(s) in name %r" % name) - - -def _read_array(stream, dtype, n): - try: - size = int(_np.dtype(dtype).itemsize * n) - return _np.frombuffer(stream.read(size), dtype) - except Exception: - raise StopIteration - - -def _write_array(stream, array): - stream.write(array.tostring()) - - -def _can_mmap(stream): - try: - pos = stream.tell() - try: - _np.memmap(stream, 'u1', 'c') - stream.seek(pos) - return True - except Exception as e: - stream.seek(pos) - return False - except Exception as e: - return False - diff --git a/dataset_generation/sort_sequences.py b/dataset_generation/sort_sequences.py deleted file mode 100755 index 1275942..0000000 --- a/dataset_generation/sort_sequences.py +++ /dev/null @@ -1,84 +0,0 @@ -#!/usr/bin/env python -from __future__ import print_function -import os -import rospy -import argparse - -import numpy as np -import rosbag -from pathlib import Path -from tqdm import tqdm -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext - - -def get_seqs(args, stamp_range=None): - if stamp_range: - pass - else: - pass - -def stamp_from_filename(fname): - stem = Path(fname).stem - msec_idx = stem.rfind('_') - msecs = float(stem[msec_idx+1:]) - print('msecs',msecs) - sec_idx = stem[:msec_idx].rfind('-') - secs = float(stem[sec_idx+1:msec_idx]) - print('secs',secs) - stamp = secs + (msecs / 1000.0) - print('stamp',stamp) - return stamp - -def seq_in_bag(seq, bag): - bag_start = bag.get_start_time() - bag_end = bag.get_end_time() - return (seq[0] >= bag_start and seq[1] <= bag_end) - -def main(args): - collection_dir = Path(args.collection_dir) - seq_stamp_map = {} # sequence id: (start time, end time) - stamp_seq_map = {} - seq_dirs = [d for d in collection_dir.iterdir() if d.is_dir()] - for sd in seq_dirs: - ddl = [d for d in sd.iterdir() if d.is_dir()] - data_files = sorted([df for dd in ddl for df in dd.iterdir() - if df.suffix in args.data_extensions]) - timestamp_begin = stamp_from_filename(data_files[0]) - timestamp_end = stamp_from_filename(data_files[-1]) - print('ts_begin',timestamp_begin,'ts_end',timestamp_end) - seq_stamp_map[str(sd)] = (timestamp_begin, timestamp_end) - stamp_seq_map[timestamp_begin] = str(sd) - sorted_seqs = [stamp_seq_map[k] for k in sorted(stamp_seq_map.keys())] - sorted_stamps = [seq_stamp_map[sk] for sk in sorted_seqs] - print('All sequences in order: ',sorted_seqs) - print('All stamps in order', sorted_stamps) - print('Stamp differences', [sorted_stamps[i][0] - sorted_stamps[i-1][1] - for i in range(len(sorted_stamps)) if i>0]) - # print('Combined',zip(sorted_seqs,sorted_stamps)) - if args.bags: - bag_seqs={} - for b in args.bags: - bp = Path(b) - bag = rosbag.Bag(b) - for sid, stamps in seq_stamp_map.items(): - if seq_in_bag(stamps, bag): - bag_seqs.setdefault(bp.stem,[]).append([sid]) - print ('bag seq map', bag_seqs) - -if __name__ == '__main__': - """Sort filenames by timestamp. Filename naming convention is assumed to be - of the form -secs_msec. - """ - parser = argparse.ArgumentParser(description="Sort multiple sequences according to timestamps of their contained data") - parser.add_argument("-d", "--collection_dir", default=None, - help="Parent dir of sequences.") - # parser.add_argument("--reorder_sequences", action="store_true", - # help="Renames sequences according to their timestamp order") - parser.add_argument("-e", "--data_extensions", nargs='+',default=['.png','.jpg','.ply'], - help="set of supported extensions of raw data files") - parser.add_argument("-b", "--bags", nargs='+', default=[], - help="output the set of sequences that correspond to the input list of rosbags") - - args = parser.parse_args() - main(args) diff --git a/dataset_generation/synchronize_frames.py b/dataset_generation/synchronize_frames.py deleted file mode 100755 index b0e459c..0000000 --- a/dataset_generation/synchronize_frames.py +++ /dev/null @@ -1,200 +0,0 @@ -#!/usr/bin/env python -from __future__ import print_function -import math -import os -import sys -import shutil -import argparse -from os import listdir, makedirs -from os.path import exists, isfile, join, splitext -import rospy -from pydoc import locate -import rosbag -from message_filters import TimeSynchronizer, ApproximateTimeSynchronizer, Subscriber -from sensor_msgs.msg import Image, CameraInfo, PointCloud2 -import re -from tqdm import tqdm - -def sorted_alphanum(file_list_ordered): - convert = lambda text: int(text) if text.isdigit() else text - alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] - return sorted(file_list_ordered, key=alphanum_key) - - -def get_file_list(path, extension=None): - if extension is None: - file_list = [path + f for f in listdir(path) if isfile(join(path, f))] - else: - file_list = [path + f for f in listdir(path) - if isfile(join(path, f)) and splitext(f)[1] == extension] - file_list = sorted_alphanum(file_list) - return file_list - - -def add_if_exists(path_dataset, folder_names): - for folder_name in folder_names: - if exists(join(path_dataset, folder_name)): - path = join(path_dataset, folder_name) - return path - - -def get_rgbd_folders(path_dataset): - path_color = add_if_exists(path_dataset, ["image/", "rgb/", "color/"]) - path_depth = join(path_dataset, "depth/") - return path_color, path_depth - - -def get_rgbd_file_lists(path_dataset): - path_color, path_depth = get_rgbd_folders(path_dataset) - color_files = get_file_list(path_color, ".jpg") + \ - get_file_list(path_color, ".png") - depth_files = get_file_list(path_depth, ".png") - return color_files, depth_files - - -def make_clean_folder(path_folder): - if not exists(path_folder): - makedirs(path_folder) - else: - shutil.rmtree(path_folder) - makedirs(path_folder) - - -def check_folder_structure(path_dataset): - path_color, path_depth = get_rgbd_folders(path_dataset) - assert exists(path_depth), \ - "Path %s is not exist!" % path_depth - assert exists(path_color), \ - "Path %s is not exist!" % path_color - -def sync_cb(*args): - topics, sync_bag = args[-2:] - assert(len(topics) == len(args[:-2])) - for i,m in enumerate(args[:-2]): - sync_bag.write(topics[i], m, m.header.stamp) - -def synchronize_rosbag(args): - rospy.init_node('sync_bag_messages') - bag = rosbag.Bag(args.dataset) - sync_bag = rosbag.Bag(os.path.splitext(args.dataset)[0]+'_synced.bag', 'w') - topics = bag.get_type_and_topic_info()[1].keys() - types = [tt[0] for tt in bag.get_type_and_topic_info()[1].values()] - sync_filt_types = [(ttopic, ttype[0][ttype[0].rfind('/')+1:]) for - ttopic, ttype in bag.get_type_and_topic_info()[1].items() - if ttype in args.type_filter] - sync_filt_topics = [(ttopic, ttype[0][ttype[0].rfind('/')+1:]) for - ttopic, ttype in bag.get_type_and_topic_info()[1].items() - if ttopic in args.topic_filter] - sync_topics_types = set(sync_filt_types + sync_filt_topics) - subscribers = [Subscriber(top, getattr(sys.modules[__name__], typ)) - for top,typ in sync_topics_types] - publishers = [rospy.Publisher(top, getattr(sys.modules[__name__], typ), queue_size=100) - for top,typ in sync_topics_types] - if args.approx: - sync = ApproximateTimeSynchronizer(subscribers, queue_size=100, slop=args.slop) - else: - sync = TimeSynchronizer(subscribers, queue_size=100) - sync_topics = [top for top,typ in sync_topics_types] - sync.registerCallback(sync_cb, sync_topics, sync_bag) - rospy.sleep(0.1) - for topic, msg, t in tqdm(bag.read_messages(), total=bag.get_message_count()): - if rospy.is_shutdown(): - break - if topic in sync_topics: - idx = sync_topics.index(topic) - publishers[idx].publish(msg) - rospy.sleep(0.01) - elif args.add_remaining_msgs: - sync_bag.write(topic, msg, t) - rospy.sleep(2.0) - - print('synced', sync_bag) - - for sub in subscribers: - sub.unregister() - bag.close() - sync_bag.close() - -def synchronize_redwood(args): - folder_path = args.dataset - color_files, depth_files = get_rgbd_file_lists(folder_path) - if args.debug_mode: - print(depth_files) - print(color_files) - - # filename format is: - # frame-timestamp.filetype - timestamps = {'depth':[None] * len(depth_files), - 'color':[None] * len(color_files)} - for i, name in enumerate(depth_files): - depth_timestamp = int(os.path.basename(depth_files[i]).replace('-','.').split('.')[1]) - timestamps['depth'][i] = depth_timestamp - for i, name in enumerate(color_files): - color_timestamp = int(os.path.basename(color_files[i]).replace('-','.').split('.')[1]) - timestamps['color'][i] = color_timestamp - - # associations' index is the color frame, and the value at - # that index is the best depth frame for the color frame - associations = [] - depth_idx = 0 - for i in range(len(color_files)): - best_dist = float('inf') - while depth_idx <= len(depth_files)-1 and i <= len(color_files)-1: - dist = math.fabs(timestamps['depth'][depth_idx] - \ - timestamps['color'][i]) - if dist > best_dist: - break - best_dist = dist - depth_idx += 1 - if depth_idx > timestamps['depth'][-1]: - print("Ended at color frame %d, depth frame %d" % (i, depth_idx)) - associations.append(depth_idx-1) - if args.debug_mode: - print("%d %d %d %d" % (i, depth_idx-1, - timestamps['depth'][depth_idx-1], timestamps['color'][i])) - - os.rename(os.path.join(folder_path, "depth"), - os.path.join(folder_path, "temp")) - if not os.path.exists(os.path.join(folder_path, "depth")): - os.makedirs(os.path.join(folder_path, "depth")) - for i, assn in enumerate(associations): - temp_name = os.path.join(folder_path, "temp", - os.path.basename(depth_files[assn])) - new_name = os.path.join(folder_path, "depth/%06d-%012d.png" % (i+1, timestamps['depth'][assn])) - print("i %d, assn %d, name %s" % (i, assn, depth_files[assn])) - if args.debug_mode: - print(temp_name) - print(new_name) - if not exists(temp_name): - assert(i+1 == len(color_files)) - os.remove(color_files[-1]) - else: - os.rename(temp_name, new_name) - shutil.rmtree(os.path.join(folder_path, "temp")) - -if __name__ == '__main__': - parser = argparse.ArgumentParser(description="Synchronize topics from rosbags, or color/depth from Redwood dataset", formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("--dataset", help="path to the dataset (currently supported : Redwood dataset or rosbag)", default='redwood') - parser.add_argument("--rosbag_folder", help="path to folder of rosbags to convert", default=None) - parser.add_argument("--type_filter", help="sync messages of given types", nargs='+', default=[]) # 'sensor_msgs/CameraInfo','sensor_msgs/Image' - parser.add_argument("--topic_filter", help="sync messages of given topics", nargs='+', default=[]) - parser.add_argument("--approx", help="approximate time sync (if false, use exact time sync)", action="store_true") - parser.add_argument("--slop", help="approximate time sync slop, in seconds", type=float, default=0.05) - parser.add_argument("--debug_mode", help="turn on debug mode", action="store_true") - parser.add_argument("--add_remaining_msgs", help="preserve all messages not filtered by sync", action="store_true") - args = parser.parse_args() - - if args.rosbag_folder is not None: - bag_names = [f for f in listdir(args.rosbag_folder) if isfile(join(args.rosbag_folder, f))] - for bag_name in bag_names: - if bag_name[-4:] != ".bag" or bag_name[-11:] == '_synced.bag': - continue - args.dataset = join(args.rosbag_folder, bag_name) - print("Synchronizing:", args.dataset) - synchronize_rosbag(args) - elif args.dataset.endswith('bag'): - synchronize_rosbag(args) - elif args.dataset == 'redwood': - synchronize_redwood(args) - else: - print("Error: Unknown dataset type, ", args.dataset) diff --git a/dataset_generation/write_tfs.py b/dataset_generation/write_tfs.py deleted file mode 100644 index 57d7880..0000000 --- a/dataset_generation/write_tfs.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import print_function -import rospy -import rosbag -import argparse -from pathlib import Path -import numpy as np -import tf2_ros -from tf2_msgs.msg import TFMessage -from geometry_msgs.msg import TransformStamped -import yaml -from tqdm import tqdm - -def main(args): - root_path = Path(args.root_dir) - collection_path = root_path / Path(args.collection_name) - - bag_seq_map = {} - bag_seq_path = collection_path / Path('bag_sequence_map.yaml') - with open(str(bag_seq_path)) as f: - bag_seq_map = yaml.load(f.read()) - - tf_buff = tf2_ros.Buffer() - for bagname, seq in bag_seq_map.items(): - tf_buff.clear() - bag_path = collection_path / Path('bags') / Path(bagname) - info_files = [s for s in (collection_path / Path(str(seq).zfill(5))).iterdir() if 'info.yaml' in s.name] - sensor_names = [f.stem[:f.stem.find('_info')] for f in info_files] - print('sensor_names', sensor_names) - frame_ids = [] - for info_file in info_files: - with info_file.open('r') as f: - frame_id = yaml.load(f.read())['frame_id'] - frame_ids.append(frame_id) - print('frame_ids',frame_ids) - b = rosbag.Bag(str(bag_path)) - for topic, msg, t in b.read_messages(topics=['/tf','/tf_static']): - if topic == '/tf_static': - for m in msg.transforms: - tf_buff.set_transform_static(m, "default_authority") - - elif topic == '/tf': - for m in msg.transforms: - tf_buff.set_transform(m, "default_authority") - - seq_tfs = {} - for i, frame in tqdm(enumerate(frame_ids)): - for j in range(i+1,len(frame_ids)): - if tf_buff.can_transform(frame_ids[i], frame_ids[j], rospy.Time(0)): - tform_msg = TransformStamped() - tform_msg = tf_buff.lookup_transform(frame_ids[i], frame_ids[j], rospy.Time(0)) - t = tform_msg.transform.translation - q = tform_msg.transform.rotation - seq_tfs[sensor_names[i]+'-'+sensor_names[j]] = {'t':{'x':t.x,'y':t.y,'z':t.z}, - 'q':{'x':q.x,'y':q.y,'z':q.z,'w':q.w}} - tf_seq_path = collection_path / Path(str(seq).zfill(5)) / Path('transforms.yaml') - - with open(str(tf_seq_path), 'w') as f: - f.write(yaml.dump(seq_tfs)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-r", "--root_dir", default="", - help="dataset root dir") - parser.add_argument("-n", "--collection_name", default="rellis", - help="collection name") - args = parser.parse_args() - main(args)