Skip to content

Commit

Permalink
Own stuff working
Browse files Browse the repository at this point in the history
  • Loading branch information
@RobinSchmid7
RobinSchmid7 committed Aug 15, 2023
1 parent 8f7ef5a commit 4651cb9
Show file tree
Hide file tree
Showing 4 changed files with 301 additions and 175 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ wild_visual_navigation/notebooks*
results/*
assets/dino/*
assets/stego/cocostuff27_vit_base_5.ckpt
assets/weights/*
.pytest_cache/**
.vscode/**
notebooks/**
Expand Down
209 changes: 137 additions & 72 deletions scripts/dataset_generation/extract_images_and_labels.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from tf_bag import BagTfTransformer
import rospy
import rosparam
from sensor_msgs.msg import Image, CameraInfo
from sensor_msgs.msg import Image, CameraInfo, CompressedImage
import rosbag

from postprocessing_tools_ros.merging import merge_bags_single, merge_bags_all
Expand Down Expand Up @@ -64,29 +64,63 @@ def do(n, dry_run):

s = os.path.join(ROOT_DIR, d["name"])

valid_topics = ["/state_estimator/anymal_state", "/alphasense_driver_ros/cam4", "/log/state/desiredRobotTwist"]
# For South Africa
# valid_topics = ["/state_estimator/anymal_state", "/alphasense_driver_ros/cam4/debayered/compressed",
# "/twist_mux/twist"]
# valid_topics = ["/state_estimator/anymal_state", "/alphasense_driver_ros/cam4/debayered",
# "/twist_mux/twist"]
# valid_topics = ["/state_estimator/anymal_state", "/alphasense_driver_ros/cam4/debayered"]
# valid_topics = ["/state_estimator/anymal_state", "/alphasense_driver_ros/cam4/debayered"]

valid_topics = ["/state_estimator/anymal_state", "/wide_angle_camera_front/img_out"]

# valid_topics = ["/state_estimator/anymal_state", "/alphasense_driver_ros/cam4/dropped/debayered/slow"]

# Merge rosbags if necessary
rosbags = [
str(s)
for s in Path(s).rglob("*.bag")
if str(s).find("lpc_robot_state") != -1
or str(s).find("jetson_images") != -1
or str(s).find("lpc_locomotion") != -1
]
try:
rosbags.sort(key=lambda x: int(x.split("/")[-1][-5]))
except:
pass
# rosbags = [
# str(s)
# for s in Path(s).rglob("*.bag")
# if str(s).find("lpc_robot_state") != -1 # Contains /state_estimator/anymal_state
# or str(s).find("jetson_images") != -1 # Contains /alphasense_driver_ros/cam4/debayered/compressed
# # or str(s).find("lpc_locomotion") != -1 # Contains /twist_mux/twist
# ]
# try:
# rosbags.sort(key=lambda x: int(x.split("/")[-1][-5]))
# except:
# pass

# rosbags = ["/home/rschmid/RosBags/perugia_grass_bags/raw/lpc_robot_state_0.bag",
# "/home/rschmid/RosBags/perugia_grass_bags/raw/lpc_robot_state_1.bag",
# "/home/rschmid/RosBags/perugia_grass_bags/raw/lpc_robot_state_2.bag",
# "/home/rschmid/RosBags/perugia_grass_bags/raw/jetson_images.bag"]
# rosbags = ["/home/rschmid/RosBags/perugia/lpc_robot_state_0.bag",
# "/home/rschmid/RosBags/perugia/jetson_images_0.bag"]
rosbags = ["/home/rschmid/RosBags/6/images.bag",
"/home/rschmid/RosBags/6/2023-03-02-11-13-08_anymal-d020-lpc_mission_0.bag",
"/home/rschmid/RosBags/6/2023-03-02-11-13-08_anymal-d020-lpc_mission_1.bag"]
# rosbags = ["/home/rschmid/RosBags/4/images.bag",
# "/home/rschmid/RosBags/4/2023-03-02-10-11-56_anymal-d020-lpc_mission_0.bag",
# "/home/rschmid/RosBags/4/2023-03-02-10-11-56_anymal-d020-lpc_mission_1.bag"]
# rosbags = ["/home/rschmid/RosBags/11/images.bag",
# "/home/rschmid/RosBags/11/2023-03-02-12-31-23_anymal-d020-lpc_mission_0.bag",
# "/home/rschmid/RosBags/11/2023-03-02-12-31-23_anymal-d020-lpc_mission_1.bag",
# "/home/rschmid/RosBags/11/2023-03-02-12-31-23_anymal-d020-lpc_mission_2.bag"]

# rosbags = ["/home/rschmid/RosBags/hoengg_extract/2021-12-06-15-31-17.bag",
# "/home/rschmid/RosBags/hoengg_extract/output_2023-02-10-15-20-27.bag"]

output_bag_wvn = s + "_wvn.bag"
output_bag_tf = s + "_tf.bag"
tf_bags = [b for b in rosbags if b.find("lpc_robot_state") != -1]
# tf_bags = [b for b in rosbags if b.find("lpc_robot_state") != -1]

# jetson locomotion robot
# if not os.path.exists(output_bag_tf):
# total_included_count, total_skipped_count = merge_bags_single(
# input_bag=tf_bags, output_bag=output_bag_tf, topics="/tf /tf_static", verbose=True
# )
if not os.path.exists(output_bag_tf):
total_included_count, total_skipped_count = merge_bags_single(
input_bag=tf_bags, output_bag=output_bag_tf, topics="/tf /tf_static", verbose=True
input_bag=rosbags, output_bag=output_bag_tf, topics="/tf /tf_static", verbose=True
)
if not os.path.exists(output_bag_wvn):
total_included_count, total_skipped_count = merge_bags_single(
Expand All @@ -99,21 +133,22 @@ def do(n, dry_run):
mission = s.split("/")[-1]

# 2022-05-12T11:56:13_mission_0_day_3
# extraction_store_folder = f"/media/Data/Datasets/2022_Perugia/wvn_output/day3/{mission}"
extraction_store_folder = f"/media/matias/datasets/2022_Perugia/wvn_output/day3/{mission}"
# running_store_folder = f"/media/Data/Datasets/2022_Perugia/wvn_output/day3/{mission}"
running_store_folder = f"/home/rschmid/RosBags/output/{mission}"

if os.path.exists(extraction_store_folder):
print(f"Stopped because folder already exists: {extraction_store_folder}")
return
if os.path.exists(running_store_folder):
print("Folder already exists, but proceeding!")
# return

rosparam.set_param("wild_visual_navigation_node/mode", "extract_labels")
rosparam.set_param("wild_visual_navigation_node/extraction_store_folder", extraction_store_folder)
# rosparam.set_param("wild_visual_navigation_node/running_store_folder", running_store_folder)
rosparam.set_param("wild_visual_navigation_node/running_store_folder", running_store_folder)

# for proprioceptive callback
state_msg_valid = False
desired_twist_msg_valid = False

pub = rospy.Publisher("/alphasense_driver_ros/cam4", Image, queue_size=1)
# pub = rospy.Publisher("/wide_angle_camera_front/image_color/compressed", CompressedImage, queue_size=1)
wvn_ros_interface = WvnRosInterface()
print("-" * 80)

Expand All @@ -122,27 +157,46 @@ def do(n, dry_run):
wvn_ros_interface.setup_rosbag_replay(tf_listener)
print("done loading tf")

# Höngg new
info_msg = CameraInfo()
info_msg.height = 540
info_msg.width = 720
info_msg.distortion_model = "plumb_bob"
info_msg.K = [347.548139773951, 0.0, 342.454373227748, 0.0, 347.434712422309, 271.368057185649, 0.0, 0.0, 1.0]
info_msg.P = [
347.548139773951,
0.0,
342.454373227748,
0.0,
0.0,
347.434712422309,
271.368057185649,
0.0,
0.0,
0.0,
1.0,
0.0,
]
info_msg.height = 1080
info_msg.width = 1440
info_msg.distortion_model = "equidistant"
info_msg.D = [0.4316922809468283, 0.09279900476637248, -0.4010909691803734, 0.4756163338479413]
info_msg.K = [575.6050407221768, 0.0, 745.7312198525915, 0.0, 578.564849365178, 519.5207040671075, 0.0, 0.0, 1.0]
info_msg.P = [575.6050407221768, 0.0, 745.7312198525915, 0.0, 0.0, 578.564849365178, 519.5207040671075, 0.0, 0.0, 0.0, 1.0, 0.0]
info_msg.R = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]

# Perugia
# info_msg = CameraInfo()
# info_msg.height = 540
# info_msg.width = 720
# info_msg.distortion_model = "plumb_bob"
# info_msg.K = [347.548139773951, 0.0, 342.454373227748, 0.0, 347.434712422309, 271.368057185649, 0.0, 0.0, 1.0]
# info_msg.P = [347.548139773951, 0.0, 342.454373227748, 0.0, 0.0, 347.434712422309, 271.368057185649, 0.0,
# 0.0, 0.0, 1.0, 0.0]
# info_msg.R = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]

# Höngg
# info_msg = CameraInfo()
# info_msg.height = 540
# info_msg.width = 720
# info_msg.distortion_model = "equidistant"
# info_msg.K = [349.5636550689, 0.0, 357.2746308879, 0.0, 349.4046775293, 264.3108985411, 0.0, 0.0, 1.0]
# info_msg.P = [349.5636550689, 0.0, 357.2746308879, 0.0, 0.0, 349.4046775293, 264.3108985411, 0.0, 0.0, 0.0, 1.0, 0.0]
# info_msg.R = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]

# South Africa
# info_msg = CameraInfo()
# info_msg.height = 1080
# info_msg.width = 1440
# info_msg.distortion_model = "equidistant"
# info_msg.D = [-0.0480706813, 0.0129997684, -0.0112199955, 0.0026955514]
# info_msg.K = [699.2284099702, 0.0, 711.8009584441, 0.0, 698.546880367, 524.7993478318, 0.0, 0.0, 1.0]
# info_msg.P = [699.2284099702, 0.0, 711.8009584441, 0.0, 0.0, 698.546880367, 524.7993478318, 0.0, 0.0,
# 0.0, 1.0, 0.0]
# info_msg.R = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]

rosbag_info_dict = get_bag_info(output_bag_wvn)
total_msgs = sum([x["messages"] for x in rosbag_info_dict["topics"] if x["topic"] in valid_topics])
total_time_img = 0
Expand All @@ -160,54 +214,65 @@ def do(n, dry_run):
bar_format="{desc:<13}{percentage:3.0f}%|{bar:20}{r_bar}",
) as pbar:
for (topic, msg, ts) in bag.read_messages(topics=None, start_time=start_time, end_time=end_time):

if rospy.is_shutdown():
return
pbar.update(1)
st = time.time()
# print(topic)
if topic == "/state_estimator/anymal_state":
# print("Received /state_estimator/anymal_state")
state_msg = anymal_msg_callback(msg, return_msg=True)
state_msg_valid = True

elif topic == "/log/state/desiredRobotTwist":
desired_twist_msg = msg
desired_twist_msg_valid = True

elif topic == "/alphasense_driver_ros/cam4":
N = 1000
for i in range(N):
pub.publish(msg)
# Change this for service call
try:
image_msg = rospy.wait_for_message(
"/alphasense_driver_ros/cam4/debayered", Image, timeout=0.01
)
suc = True
except Exception as e:
suc = False
pass
if suc:
if msg.header.stamp == image_msg.header.stamp:
break
if i >= N - 1:
raise Exception("Timeout waiting for debayered image message")
# elif topic == "/twist_mux/twist":
# # print("Received /twist_mux/twist")
# desired_twist_msg = msg
# desired_twist_msg_valid = True

elif topic == "/wide_angle_camera_front/img_out":
image_msg = msg
# Image is already debayered, need to recompress it
print("Received /wide_angle_camera_front/img_out")
# for i in range(100):
# pub.publish(msg)
# try:
# image_msg = rospy.wait_for_message(
# "/wide_angle_camera_front/image_color", Image, timeout=0.1
# )
# suc = True
# except:
# suc = False
# pass
# if suc:
# if msg.header.stamp == image_msg.header.stamp:
# break
# if i >= 99:
# raise Exception("Timeout waiting for debayered image message")

info_msg.header = msg.header
camera_options = {}
camera_options['name'] = "wide_angle_camera_front"
camera_options["use_for_training"] = True

try:
wvn_ros_interface.image_callback(image_msg, info_msg)
wvn_ros_interface.image_callback(image_msg, info_msg, camera_options)
except Exception as e:
tqdm.write("Bad image_callback", e)
print("Bad image_callback", e)

total_time_img += time.time() - st
# print(f"image time: {total_time_img} , state time: {total_time_state}")
tqdm.write("add image")
if state_msg_valid and desired_twist_msg_valid:
print("add image")
if state_msg_valid:
try:
wvn_ros_interface.robot_state_callback(state_msg, desired_twist_msg)
wvn_ros_interface.robot_state_callback(state_msg, None)
except Exception as e:
tqdm.write("Bad robot_state callback ", e)
print("Bad robot_state callback ", e)

state_msg_valid = False
desired_twist_msg_valid = True
# desired_twist_msg_valid = True
total_time_state += time.time() - st
tqdm.write("add supervision")
# print("add supervision")

print("Finished with converting the dataset")
rospy.signal_shutdown("stop the node")
Expand All @@ -217,8 +282,8 @@ def do(n, dry_run):
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--n", type=int, default=2, help="Store data")
parser.add_argument("--n", type=int, default=0, help="Store data")
parser.add_argument("--dry_run", type=int, default=0, help="Store data")
args = parser.parse_args()
print(args.n)

do(args.n, args.dry_run)
Loading

0 comments on commit 4651cb9

Please sign in to comment.