Skip to content

Commit

Permalink
refactored code, reformatted all visual elements into new ShapeScene …
Browse files Browse the repository at this point in the history 
…class
  • Loading branch information
@K0-p
K0-p committed Aug 22, 2024
1 parent c43d623 commit 59b9dc8
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 210 deletions.
159 changes: 53 additions & 106 deletions UI/foxglove/UI_launcher.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,7 @@

# local tools to manage Foxglove scenes
from plot.foxglove_utils import (
FoxgloveShapeListener,
SceneChannel,
ScalableArrowsScene,
ShapeScene
)

Expand Down Expand Up @@ -140,20 +138,6 @@ async def main():
SceneUpdate.DESCRIPTOR.full_name,
scene_schema,
).add_chan(server)
test_chan_id = await SceneChannel(
False,
"n_test_viz",
"protobuf",
SceneUpdate.DESCRIPTOR.full_name,
scene_schema,
).add_chan(server)
sphere_test_chan_id = await SceneChannel(
False,
"00_spheres",
"protobuf",
SceneUpdate.DESCRIPTOR.full_name,
scene_schema,
).add_chan(server)
grfs_chan_id = await SceneChannel(
True,
"GRFs",
Expand All @@ -170,7 +154,6 @@ async def main():
"rfoot_rf_normal_filt",
],
).add_chan(server)

icpS_chan_id = await SceneChannel(
False, "icp_viz", "protobuf", SceneUpdate.DESCRIPTOR.full_name, scene_schema
).add_chan(server)
Expand All @@ -181,24 +164,15 @@ async def main():
for scn in range(len(xyz_scene_names)):
await sceneinitman(xyz_scene_names[scn], server)

# create all of the visual scenes
scenes = []
icp_listener = FoxgloveShapeListener(
icpS_chan_id,
"spheres",
["est_icp", "des_icp"],
{"est_icp": [0.1, 0.1, 0.1], "des_icp": [0.1, 0.1, 0.1]},
{"est_icp": [1, 0, 1, 1], "des_icp": [0, 1, 0, 1]},
)
scenes.append(icp_listener)

# STEP_MAX available footstep plans
hfoot_length, hfoot_width = foot_dimensions()
x = 2*hfoot_length
y = 2*hfoot_width
msgs = []
proj_footstep_chan_ids = []
proj_footstep_viz_chan_ids = []
proj_foot_pos, proj_foot_ori = {}, {} # Stores step position
proj_feet = []
for i in range(STEP_MAX):
rf = "proj_rf"+str(i)
lf = "proj_lf"+str(i)
Expand All @@ -220,26 +194,23 @@ async def main():
proj_footstep_viz_chan_id = await SceneChannel(
False, name+"_viz", "protobuf", SceneUpdate.DESCRIPTOR.full_name, scene_schema
).add_chan(server)
proj_footstep_listener = FoxgloveShapeListener(proj_footstep_viz_chan_id, "cubes",
[rf, lf], {rf: [x, y, 0.001], lf: [x, y, 0.001]}, {rf: [1, 0, 0, 1], lf: [.1, .5, 1, 1]})
scenes.append(proj_footstep_listener)

arrows_scene = ScalableArrowsScene()
arrows_scene.add_arrow("lfoot_rf_cmd", [0, 0, 1, 0.5]) # blue arrow
arrows_scene.add_arrow("rfoot_rf_cmd", [0, 0, 1, 0.5]) # blue arrow
arrows_scene.add_arrow(
"lfoot_rf_normal_filt", [0.2, 0.2, 0.2, 0.5]
) # grey arrow
arrows_scene.add_arrow(
"rfoot_rf_normal_filt", [0.2, 0.2, 0.2, 0.5]
) # grey arrow
scenecount = len(scenes) - 1

new_arrows = ScalableArrowsScene()
new_arrows.add_arrow("cheese",[0, 0, 1, 0.5])

test_sphere = ShapeScene()
test_sphere.add_shape("round", "spheres", [.1,.5,1,1])
proj_footstep_viz_chan_ids.append(proj_footstep_viz_chan_id)
#add footstep visual elements
proj_footsteps = ShapeScene()
proj_footsteps.add_shape(rf, "cubes", [1, 0, 0, 1], [x, y, 0.001])
proj_footsteps.add_shape(lf, "cubes", [.1, .5, 1, 1], [x, y, 0.001])
proj_feet.append(proj_footsteps)

arrows_scene = ShapeScene()
arrows_scene.add_shape("lfoot_rf_cmd", "arrows", [0, 0, 1, 0.5], [0.03, 0.1, 0.08]) # blue arrow
arrows_scene.add_shape("rfoot_rf_cmd", "arrows", [0, 0, 1, 0.5], [0.03, 0.1, 0.08]) # blue arrow
arrows_scene.add_shape("lfoot_rf_normal_filt", "arrows", [0.2, 0.2, 0.2, 0.5], [0.03, 0.1, 0.08]) # grey arrow
arrows_scene.add_shape("rfoot_rf_normal_filt", "arrows", [0.2, 0.2, 0.2, 0.5], [0.03, 0.1, 0.08]) # grey arrow

#add visual icp spheres
icp_spheres = ShapeScene()
icp_spheres.add_shape("est_icp", "spheres", [1, 0, 1, 1], [0.1, 0.1, 0.1])
icp_spheres.add_shape("des_icp", "spheres", [0, 1, 0, 1], [0.1, 0.1, 0.1])

# Send the FrameTransform every frame to update the model's position
transform = FrameTransform()
Expand All @@ -256,13 +227,7 @@ async def main():
observer.start()

while True:
# cycle through all visual scenes --CAUSES A BUG WHERE NOT ALL VIZ SCENES SHOW
server.set_listener(scenes[scenecount])
tasks = []

scenecount = scenecount - 1
if scenecount == -1:
scenecount = len(scenes) - 1
tasks = [] # Scenes to synchronously update

# receive msg trough socket
encoded_msg = socket.recv()
Expand Down Expand Up @@ -393,38 +358,6 @@ async def main():
transform.rotation.Clear()
transform.translation.Clear()

# update icp values on the grid
for obj in ["est_icp", "des_icp"]:
transform.parent_frame_id = "world"
transform.child_frame_id = obj
transform.timestamp.FromNanoseconds(now)
transform.translation.x = list(getattr(msg, obj))[0]
transform.translation.y = list(getattr(msg, obj))[1]
tasks.append(server.send_message(
tf_chan_id, now, transform.SerializeToString()
))

update = fp.sd.steps_to_update()
for obj in msgs:
if obj in update:
yaml = fp.sd.yaml_num
setattr(fp.sd, obj[5] + "f_steps_taken", getattr(fp.sd, obj[5] + "f_steps_taken") + 1)
proj_foot_pos[obj] = getattr(fp, obj[5] + "foot_contact_pos")[yaml][update[obj]]
proj_foot_ori[obj] = getattr(fp, obj[5] + "foot_contact_ori")[yaml][update[obj]]
transform.parent_frame_id = "world"
transform.child_frame_id = obj
transform.timestamp.FromNanoseconds(now)
transform.translation.x = proj_foot_pos[obj][0]
transform.translation.y = proj_foot_pos[obj][1]
transform.translation.z = proj_foot_pos[obj][2]
transform.rotation.x = proj_foot_ori[obj][1]
transform.rotation.y = proj_foot_ori[obj][2]
transform.rotation.z = proj_foot_ori[obj][3]
transform.rotation.w = proj_foot_ori[obj][0]
tasks.append(server.send_message(
tf_chan_id, now, transform.SerializeToString()
))

Ry = R.from_euler("y", -np.pi / 2).as_matrix()
# update GRF arrows
for obj in [
Expand Down Expand Up @@ -492,41 +425,55 @@ async def main():

# force scale
force_magnitude = force_magnitude / 1200.0
arrows_scene.update(obj, quat_force, force_magnitude, now)
arrows_scene.scale(obj, quat_force, force_magnitude, now)
tasks.append(server.send_message(
tf_chan_id, now, transform.SerializeToString()
))
tasks.append(server.send_message(
normS_chan_id, now, arrows_scene.serialized_msg(obj)
))

for obj in ["cheese"]:
# update icp values on the grid
for obj in ["est_icp", "des_icp"]:
transform.parent_frame_id = "world"
transform.child_frame_id = obj
transform.timestamp.FromNanoseconds(now)
transform.translation.x = msg.rfoot_pos[0]
transform.translation.y = msg.rfoot_pos[1]
new_arrows.update(obj, quat_force, force_magnitude, now)
transform.translation.x = list(getattr(msg, obj))[0]
transform.translation.y = list(getattr(msg, obj))[1]
icp_spheres.update(obj, now)
tasks.append(server.send_message(
tf_chan_id, now, transform.SerializeToString()
))
tasks.append(server.send_message(
test_chan_id, now, new_arrows.serialized_msg(obj)
icpS_chan_id, now, icp_spheres.serialized_msg(obj)
))

for obj in ["round"]:
transform.parent_frame_id = "world"
transform.child_frame_id = obj
transform.timestamp.FromNanoseconds(now)
transform.translation.x = msg.rfoot_pos[0]
transform.translation.y = msg.rfoot_pos[1]
test_sphere.update(obj, quat_force, force_magnitude, now)
tasks.append(server.send_message(
tf_chan_id, now, transform.SerializeToString()
))
tasks.append(server.send_message(
sphere_test_chan_id, now, test_sphere.serialized_msg(obj)
))
# Update projected footsteps
update = fp.sd.steps_to_update()
for obj in msgs:
if obj in update:
stepnum = int(''.join(filter(lambda i: i.isdigit(), obj)))
yaml = fp.sd.yaml_num
setattr(fp.sd, obj[5] + "f_steps_taken", getattr(fp.sd, obj[5] + "f_steps_taken") + 1)
proj_foot_pos[obj] = getattr(fp, obj[5] + "foot_contact_pos")[yaml][update[obj]]
proj_foot_ori[obj] = getattr(fp, obj[5] + "foot_contact_ori")[yaml][update[obj]]
transform.parent_frame_id = "world"
transform.child_frame_id = obj
transform.timestamp.FromNanoseconds(now)
transform.translation.x = proj_foot_pos[obj][0]
transform.translation.y = proj_foot_pos[obj][1]
transform.translation.z = proj_foot_pos[obj][2]
transform.rotation.x = proj_foot_ori[obj][1]
transform.rotation.y = proj_foot_ori[obj][2]
transform.rotation.z = proj_foot_ori[obj][3]
transform.rotation.w = proj_foot_ori[obj][0]
proj_feet[stepnum].update(obj, now)
tasks.append(server.send_message(
tf_chan_id, now, transform.SerializeToString()
))
tasks.append(server.send_message(
proj_footstep_viz_chan_ids[stepnum], now, proj_feet[stepnum].serialized_msg(obj)
))

await asyncio.gather(*tasks)

Expand Down
125 changes: 21 additions & 104 deletions plot/foxglove_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,50 +12,11 @@ def SubtopicGen(names):
return json.dumps({"type": "object", "properties": subs})


class ScalableArrowsScene:
def __init__(self):
self.arrows = {}

def add_arrow(self, name, color):
arrow_update = SceneUpdate()
arrow_entity = arrow_update.entities.add()
arrow_entity.id = name
arrow_entity.frame_id = name
arrow_model = arrow_entity.arrows.add()
arrow_model.color.r = color[0]
arrow_model.color.g = color[1]
arrow_model.color.b = color[2]
arrow_model.color.a = color[3]
arrow_model.shaft_diameter = 0.03
arrow_model.head_length = 0.1
arrow_model.head_diameter = 0.08
# update dictionary of arrows to visualize
self.arrows[name] = arrow_update

def update(self, name, quat_force, force_magnitude, timestamp):
# update timestamp
self.arrows[name].entities[0].timestamp.FromNanoseconds(timestamp)

# update force direction
self.arrows[name].entities[0].arrows[0].pose.orientation.x = quat_force[0]
self.arrows[name].entities[0].arrows[0].pose.orientation.y = quat_force[1]
self.arrows[name].entities[0].arrows[0].pose.orientation.z = quat_force[2]
self.arrows[name].entities[0].arrows[0].pose.orientation.w = quat_force[3]

# update force scale
self.arrows[name].entities[0].arrows[
0
].shaft_length = force_magnitude # scale down

def serialized_msg(self, name):
return self.arrows[name].SerializeToString()


class ShapeScene:
def __init__(self):
self.shapes = {}

def add_shape(self, name, shape, color):
def add_shape(self, name, shape, color, size):
shape_update = SceneUpdate()
shape_entity = shape_update.entities.add()
shape_entity.id = name
Expand All @@ -66,20 +27,32 @@ def add_shape(self, name, shape, color):
shape_model.color.g = color[1]
shape_model.color.b = color[2]
shape_model.color.a = color[3]
shape_model.size.x = 0.1
shape_model.size.y = 0.5
shape_model.size.z = 0.1
if shape == "arrows":
shape_model.shaft_diameter = size[0]
shape_model.head_length = size[1]
shape_model.head_diameter = size[2]
else:
shape_model.size.x = size[0]
shape_model.size.y = size[1]
shape_model.size.z = size[2]
# update dictionary of shapes to visualize
self.shapes[name] = shape_update

def update(self, name, quat_force, force_magnitude, timestamp):
def update(self, name, timestamp):
self.shapes[name].entities[0].timestamp.FromNanoseconds(timestamp)

def scale(self, name, quat_force, force_magnitude, timestamp):
# update timestamp
self.shapes[name].entities[0].timestamp.FromNanoseconds(timestamp)
# update force direction
#self.shapes[name].entities[0].arrows[0].pose.orientation.x = quat_force[0]
#self.shapes[name].entities[0].arrows[0].pose.orientation.y = quat_force[1]
#self.shapes[name].entities[0].arrows[0].pose.orientation.z = quat_force[2]
#self.shapes[name].entities[0].arrows[0].pose.orientation.w = quat_force[3]
self.shapes[name].entities[0].arrows[0].pose.orientation.x = quat_force[0]
self.shapes[name].entities[0].arrows[0].pose.orientation.y = quat_force[1]
self.shapes[name].entities[0].arrows[0].pose.orientation.z = quat_force[2]
self.shapes[name].entities[0].arrows[0].pose.orientation.w = quat_force[3]
# update force scale
self.shapes[name].entities[0].arrows[
0
].shaft_length = force_magnitude # scale down

def serialized_msg(self, name):
return self.shapes[name].SerializeToString()
Expand Down Expand Up @@ -114,59 +87,3 @@ def add_chan(self, server):
}
)


def FormHandler(entity, shape, size):
shape_get = getattr(entity, shape)
nodel = shape_get.add()
if shape == "spheres" or "cubes":
nodel.size.x = size[0]
nodel.size.y = size[1]
nodel.size.z = size[2]
# Currently, the arrows change length so are visualized differently
# if(shape == "arrows"):
# nodel.shaft_length = size[0]
# nodel.shaft_diameter = size[1]
# nodel.head_length = size[2]
# nodel.head_diameter = size[3]
return nodel


class FoxgloveShapeListener(FoxgloveServerListener):
def __init__(self, scene_chanel_id, shape, msgs_list, size_dict, color_dict):
self.scene_chan_id = scene_chanel_id
self.shape = shape
self.msgs_list = msgs_list
self.size_dict = size_dict
self.color_dict = color_dict

def on_subscribe(self, server, channel_id):
msgs_list = self.msgs_list
size_dict = self.size_dict
color_dict = self.color_dict
scene_chan_id = self.scene_chan_id
now = time.time_ns()
if channel_id == scene_chan_id:
scene_update = SceneUpdate()
for msg in msgs_list:
rgb = color_dict[msg]
entity = scene_update.entities.add()
entity.timestamp.FromNanoseconds(now)
entity.id = msg
entity.frame_id = msg
entity.frame_locked = True
# shape_get = getattr(entity, self.shape)
# nodel = shape_get.add()
nodel = FormHandler(entity, self.shape, size_dict[msg])
nodel.color.r = rgb[0]
nodel.color.g = rgb[1]
nodel.color.b = rgb[2]
nodel.color.a = rgb[3]

asyncio.create_task(
server.send_message(
self.scene_chan_id, now, scene_update.SerializeToString()
)
)

def on_unsubscribe(self, server, channel_id):
pass

0 comments on commit 59b9dc8

Please sign in to comment.