155 lines
5.5 KiB
Python
155 lines
5.5 KiB
Python
import lcm
|
|
import threading
|
|
import time
|
|
import select
|
|
|
|
import numpy as np
|
|
|
|
from go1_gym_deploy.lcm_types.leg_control_data_lcmt import leg_control_data_lcmt
|
|
from go1_gym_deploy.lcm_types.rc_command_lcmt import rc_command_lcmt
|
|
from go1_gym_deploy.lcm_types.state_estimator_lcmt import state_estimator_lcmt
|
|
from go1_gym_deploy.lcm_types.vicon_pose_lcmt import vicon_pose_lcmt
|
|
from go1_gym_deploy.lcm_types.camera_message_lcmt import camera_message_lcmt
|
|
from go1_gym_deploy.lcm_types.camera_message_rect_wide import camera_message_rect_wide
|
|
from go1_gym_deploy.lcm_types.camera_message_rect_wide_mask import camera_message_rect_wide_mask
|
|
|
|
|
|
class UnitreeLCMInspector:
|
|
def __init__(self, lc):
|
|
self.lc = lc
|
|
|
|
self.camera_names = ["front", "bottom", "left", "right", "rear"]
|
|
for cam_name in self.camera_names:
|
|
self.camera_subscription = self.lc.subscribe(f"rect_image_{cam_name}", self._rect_camera_cb)
|
|
for cam_name in self.camera_names:
|
|
self.camera_subscription = self.lc.subscribe(f"rect_image_{cam_name}_mask", self._mask_camera_cb)
|
|
|
|
self.camera_image_left = None
|
|
self.camera_image_right = None
|
|
self.camera_image_front = None
|
|
self.camera_image_bottom = None
|
|
self.camera_image_rear = None
|
|
|
|
self.ts = [time.time(), time.time(), time.time(), time.time(), time.time(),]
|
|
|
|
self.num_low_states = 0
|
|
|
|
def _rect_camera_cb(self, channel, data):
|
|
|
|
# message_types = [camera_message_rect_front, camera_message_rect_front_chin, camera_message_rect_left,
|
|
# camera_message_rect_right, camera_message_rect_rear_down]
|
|
# image_shapes = [(200, 200, 3), (100, 100, 3), (100, 232, 3), (100, 232, 3), (200, 200, 3)]
|
|
|
|
message_types = [camera_message_rect_wide, camera_message_rect_wide, camera_message_rect_wide, camera_message_rect_wide, camera_message_rect_wide]
|
|
image_shapes = [(116, 100, 3), (116, 100, 3), (116, 100, 3), (116, 100, 3), (116, 100, 3)]
|
|
|
|
cam_name = channel.split("_")[-1]
|
|
|
|
cam_id = self.camera_names.index(cam_name) + 1
|
|
|
|
print(channel, message_types[cam_id - 1])
|
|
msg = message_types[cam_id - 1].decode(data)
|
|
|
|
img = np.fromstring(msg.data, dtype=np.uint8)
|
|
img = np.flip(np.flip(
|
|
img.reshape((image_shapes[cam_id - 1][2], image_shapes[cam_id - 1][1], image_shapes[cam_id - 1][0])),
|
|
axis=0), axis=1).transpose(1, 2, 0)
|
|
|
|
if cam_id == 1:
|
|
self.camera_image_front = img
|
|
elif cam_id == 2:
|
|
self.camera_image_bottom = img
|
|
elif cam_id == 3:
|
|
self.camera_image_left = img
|
|
elif cam_id == 4:
|
|
self.camera_image_right = img
|
|
elif cam_id == 5:
|
|
self.camera_image_rear = img
|
|
else:
|
|
print("Image received from camera with unknown ID#!")
|
|
|
|
print(f"f{1. / (time.time() - self.ts[cam_id - 1])}: received py from {cam_name}!")
|
|
self.ts[cam_id-1] = time.time()
|
|
|
|
from PIL import Image
|
|
im = Image.fromarray(img)
|
|
im.save(f"{cam_name}_image.jpeg")
|
|
|
|
def _mask_camera_cb(self, channel, data):
|
|
|
|
message_types = [camera_message_rect_wide_mask for i in range(5)]
|
|
image_shapes = [(116, 100, 1) for i in range(5)]
|
|
|
|
cam_name = channel.split("_")[-2]
|
|
|
|
cam_id = self.camera_names.index(cam_name) + 1
|
|
|
|
print(channel, message_types[cam_id - 1])
|
|
msg = message_types[cam_id - 1].decode(data)
|
|
|
|
img = np.array(list(msg.data)).astype(np.uint8)
|
|
img = np.flip(np.flip(
|
|
img.reshape((image_shapes[cam_id - 1][2], image_shapes[cam_id - 1][1], image_shapes[cam_id - 1][0])),
|
|
axis=0), axis=1)
|
|
|
|
if cam_id == 1:
|
|
self.camera_image_front = img
|
|
elif cam_id == 2:
|
|
self.camera_image_bottom = img
|
|
elif cam_id == 3:
|
|
self.camera_image_left = img
|
|
elif cam_id == 4:
|
|
self.camera_image_right = img
|
|
elif cam_id == 5:
|
|
self.camera_image_rear = img
|
|
else:
|
|
print("Image received from camera with unknown ID#!")
|
|
|
|
print(f"f{1. / (time.time() - self.ts[cam_id - 1])}: received py from {cam_name}!")
|
|
self.ts[cam_id-1] = time.time()
|
|
|
|
from PIL import Image
|
|
# print(img[0])
|
|
im = Image.fromarray(img[0])
|
|
im.save(f"{cam_name}_mask_image.jpeg")
|
|
|
|
def publish_30Hz(self):
|
|
msg = camera_message_rect_wide()
|
|
msg.data = [0] * 34800
|
|
self.lc.publish("rect_image_rear", msg.encode())
|
|
time.sleep(1./30.)
|
|
|
|
def poll(self, cb=None):
|
|
t = time.time()
|
|
try:
|
|
while True:
|
|
timeout = 0.01
|
|
rfds, wfds, efds = select.select([self.lc.fileno()], [], [], timeout)
|
|
if rfds:
|
|
# print("message received!")
|
|
self.lc.handle()
|
|
# print(f'Freq {1. / (time.time() - t)} Hz'); t = time.time()
|
|
else:
|
|
continue
|
|
# print(f'waiting for message... Freq {1. / (time.time() - t)} Hz'); t = time.time()
|
|
# if cb is not None:
|
|
# cb()
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
def spin(self):
|
|
self.run_thread = threading.Thread(target=self.poll, daemon=False)
|
|
self.run_thread.start()
|
|
|
|
if __name__ == "__main__":
|
|
import lcm
|
|
|
|
lc = lcm.LCM("udpm://239.255.76.67:7667?ttl=255")
|
|
#check_lcm_msgs()
|
|
print("init")
|
|
insp = UnitreeLCMInspector(lc)
|
|
print("polling")
|
|
# insp.poll()
|
|
while True:
|
|
insp.publish_30Hz()
|