silencht
1 year ago
committed by
Your Name
Your Name
7 changed files with 726 additions and 143 deletions
-
1.gitignore
-
2teleop/image_server/image_client.py
-
215teleop/image_server/image_server.py
-
219teleop/robot_control/robot_hand_unitree.py
-
76teleop/teleop_hand_and_arm.py
-
13teleop/utils/episode_writer.py
-
231teleop/utils/rerun_visualizer.py
@ -0,0 +1,231 @@ |
|||
import os |
|||
import json |
|||
import cv2 |
|||
import time |
|||
import rerun as rr |
|||
import rerun.blueprint as rrb |
|||
from datetime import datetime |
|||
|
|||
class RerunEpisodeReader: |
|||
def __init__(self, task_dir = ".", json_file="data.json"): |
|||
self.task_dir = task_dir |
|||
self.json_file = json_file |
|||
|
|||
def return_episode_data(self, episode_idx): |
|||
# Load episode data on-demand |
|||
episode_dir = os.path.join(self.task_dir, f"episode_{episode_idx:04d}") |
|||
json_path = os.path.join(episode_dir, self.json_file) |
|||
|
|||
if not os.path.exists(json_path): |
|||
raise FileNotFoundError(f"Episode {episode_idx} data.json not found.") |
|||
|
|||
with open(json_path, 'r', encoding='utf-8') as jsonf: |
|||
json_file = json.load(jsonf) |
|||
|
|||
episode_data = [] |
|||
|
|||
# Loop over the data entries and process each one |
|||
for item_data in json_file['data']: |
|||
# Process images and other data |
|||
colors = self._process_images(item_data, 'colors', episode_dir) |
|||
depths = self._process_images(item_data, 'depths', episode_dir) |
|||
audios = self._process_audio(item_data, 'audios', episode_dir) |
|||
|
|||
# Append the data in the item_data list |
|||
episode_data.append( |
|||
{ |
|||
'idx': item_data.get('idx', 0), |
|||
'colors': colors, |
|||
'depths': depths, |
|||
'states': item_data.get('states', {}), |
|||
'actions': item_data.get('actions', {}), |
|||
'tactiles': item_data.get('tactiles', {}), |
|||
'audios': audios, |
|||
} |
|||
) |
|||
|
|||
return episode_data |
|||
|
|||
def _process_images(self, item_data, data_type, dir_path): |
|||
images = {} |
|||
|
|||
for key, file_name in item_data.get(data_type, {}).items(): |
|||
if file_name: |
|||
file_path = os.path.join(dir_path, file_name) |
|||
if os.path.exists(file_path): |
|||
image = cv2.imread(file_path) |
|||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|||
images[key] = image |
|||
return images |
|||
|
|||
def _process_audio(self, item_data, data_type, episode_dir): |
|||
audio_data = {} |
|||
dir_path = os.path.join(episode_dir, data_type) |
|||
|
|||
for key, file_name in item_data.get(data_type, {}).items(): |
|||
if file_name: |
|||
file_path = os.path.join(dir_path, file_name) |
|||
if os.path.exists(file_path): |
|||
pass # Handle audio data if needed |
|||
return audio_data |
|||
|
|||
class RerunLogger: |
|||
def __init__(self, prefix = "", IdxRangeBoundary = 30, memory_limit = "200MB"): |
|||
rr.init(datetime.now().strftime("Runtime_%Y%m%d_%H%M%S")) |
|||
rr.spawn(memory_limit = memory_limit) |
|||
|
|||
self.prefix = prefix |
|||
self.IdxRangeBoundary = IdxRangeBoundary |
|||
# Set up blueprint for live visualization |
|||
if self.IdxRangeBoundary: |
|||
self.setup_blueprint() |
|||
|
|||
def setup_blueprint(self): |
|||
data_plot_paths = [ |
|||
f"{self.prefix}left_arm", |
|||
f"{self.prefix}right_arm", |
|||
f"{self.prefix}left_hand", |
|||
f"{self.prefix}right_hand" |
|||
] |
|||
image_plot_paths = [ |
|||
f"{self.prefix}colors/color_0", |
|||
f"{self.prefix}colors/color_1", |
|||
f"{self.prefix}colors/color_2", |
|||
f"{self.prefix}colors/color_3" |
|||
] |
|||
|
|||
views = [] |
|||
for plot_path in data_plot_paths: |
|||
view = rrb.TimeSeriesView( |
|||
origin = plot_path, |
|||
time_ranges=[ |
|||
rrb.VisibleTimeRange( |
|||
"idx", |
|||
start = rrb.TimeRangeBoundary.cursor_relative(seq = -self.IdxRangeBoundary), |
|||
end = rrb.TimeRangeBoundary.cursor_relative(), |
|||
) |
|||
], |
|||
plot_legend = rrb.PlotLegend(visible = True), |
|||
) |
|||
views.append(view) |
|||
|
|||
for plot_path in image_plot_paths: |
|||
view = rrb.Spatial2DView( |
|||
origin = plot_path, |
|||
time_ranges=[ |
|||
rrb.VisibleTimeRange( |
|||
"idx", |
|||
start = rrb.TimeRangeBoundary.cursor_relative(seq = -self.IdxRangeBoundary), |
|||
end = rrb.TimeRangeBoundary.cursor_relative(), |
|||
) |
|||
], |
|||
) |
|||
views.append(view) |
|||
|
|||
grid = rrb.Grid(contents = views, |
|||
grid_columns=2, |
|||
column_shares=[1, 1], |
|||
row_shares=[1, 1, 0.5], |
|||
) |
|||
rr.send_blueprint(grid) |
|||
|
|||
|
|||
def log_item_data(self, item_data: dict): |
|||
rr.set_time_sequence("idx", item_data.get('idx', 0)) |
|||
|
|||
# Log states |
|||
states = item_data.get('states', {}) or {} |
|||
for part, state_info in states.items(): |
|||
if state_info: |
|||
values = state_info.get('qpos', []) |
|||
for idx, val in enumerate(values): |
|||
rr.log(f"{self.prefix}{part}/states/qpos/{idx}", rr.Scalar(val)) |
|||
|
|||
# Log actions |
|||
actions = item_data.get('actions', {}) or {} |
|||
for part, action_info in actions.items(): |
|||
if action_info: |
|||
values = action_info.get('qpos', []) |
|||
for idx, val in enumerate(values): |
|||
rr.log(f"{self.prefix}{part}/actions/qpos/{idx}", rr.Scalar(val)) |
|||
|
|||
# Log colors (images) |
|||
colors = item_data.get('colors', {}) or {} |
|||
for color_key, color_val in colors.items(): |
|||
if color_val is not None: |
|||
rr.log(f"{self.prefix}colors/{color_key}", rr.Image(color_val)) |
|||
|
|||
# Log depths (images) |
|||
depths = item_data.get('depths', {}) or {} |
|||
for depth_key, depth_val in depths.items(): |
|||
if depth_val is not None: |
|||
# rr.log(f"{self.prefix}depths/{depth_key}", rr.Image(depth_val)) |
|||
pass # Handle depth if needed |
|||
|
|||
# Log tactile if needed |
|||
tactiles = item_data.get('tactiles', {}) or {} |
|||
for hand, tactile_vals in tactiles.items(): |
|||
if tactile_vals is not None: |
|||
pass # Handle tactile if needed |
|||
|
|||
# Log audios if needed |
|||
audios = item_data.get('audios', {}) or {} |
|||
for audio_key, audio_val in audios.items(): |
|||
if audio_val is not None: |
|||
pass # Handle audios if needed |
|||
|
|||
def log_episode_data(self, episode_data: list): |
|||
for item_data in episode_data: |
|||
self.log_item_data(item_data) |
|||
|
|||
|
|||
if __name__ == "__main__": |
|||
import gdown |
|||
import zipfile |
|||
import os |
|||
url = "https://drive.google.com/file/d/1f5UuFl1z_gaByg_7jDRj1_NxfJZh2evD/view?usp=sharing" |
|||
zip_file = "rerun_testdata.zip" |
|||
if not os.path.exists("episode_0006"): |
|||
if not os.path.exists(zip_file): |
|||
file_id = url.split('/')[5] |
|||
gdown.download(id=file_id, output=zip_file, quiet=False) |
|||
print("download ok.") |
|||
with zipfile.ZipFile(zip_file, 'r') as zip_ref: |
|||
zip_ref.extractall(".") |
|||
print("uncompress ok.") |
|||
os.remove(zip_file) |
|||
print("clean file ok.") |
|||
else: |
|||
print("rerun_testdata exits.") |
|||
|
|||
episode_reader = RerunEpisodeReader() |
|||
episode_data6 = episode_reader.return_episode_data(6) |
|||
episode_data8 = episode_reader.return_episode_data(8) |
|||
|
|||
# Example 1: Offline Visualization |
|||
user_input = input("Please enter the start signal (enter 'off' to start the subsequent program):\n") |
|||
if user_input.lower() == 'off': |
|||
print("Starting offline visualization...") |
|||
offline_logger = RerunLogger(prefix="offline/") |
|||
offline_logger.log_episode_data(episode_data6) |
|||
print("Offline visualization completed.") |
|||
|
|||
# Example 2: Online Visualization with Fixed Time Window |
|||
user_input = input("Please enter the start signal (enter 'on' to start the subsequent program):\n") |
|||
if user_input.lower() == 'on': |
|||
print("Starting online visualization with fixed idx size...") |
|||
online_logger = RerunLogger(prefix="online/", IdxRangeBoundary = 60) |
|||
for item_data in episode_data6: |
|||
online_logger.log_item_data(item_data) |
|||
time.sleep(0.033) # 30hz |
|||
print("Online visualization completed.") |
|||
|
|||
# Example 3: Online Visualization with Fixed Time Window |
|||
user_input = input("Please enter the start signal (enter 'on' to start the subsequent program):\n") |
|||
if user_input.lower() == 'on': |
|||
print("Starting online visualization with fixed idx size...") |
|||
online_logger = RerunLogger(prefix="online/", IdxRangeBoundary = 60) |
|||
for item_data in episode_data8: |
|||
online_logger.log_item_data(item_data) |
|||
time.sleep(0.033) # 30hz |
|||
print("Online visualization completed.") |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue