wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix

This commit is contained in:
Remi Cadene 2024-07-18 18:48:41 +02:00
parent 895182b272
commit 1664dac300
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lerobot/scripts/control_robot.py
View File

@ -313,7 +313,7 @@ def record_dataset(
if not is_headless(): if not is_headless():
image_keys = [key for key in observation if "image" in key] image_keys = [key for key in observation if "image" in key]
for key in image_keys: for key in image_keys:
cv2.imshow(key, convert_torch_image_to_cv2(observation[key])) cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
cv2.waitKey(1) cv2.waitKey(1)
dt_s = time.perf_counter() - now dt_s = time.perf_counter() - now
@ -354,7 +354,7 @@ def record_dataset(
if not is_headless(): if not is_headless():
image_keys = [key for key in observation if "image" in key] image_keys = [key for key in observation if "image" in key]
for key in image_keys: for key in image_keys:
cv2.imshow(key, convert_torch_image_to_cv2(observation[key])) cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
cv2.waitKey(1) cv2.waitKey(1)
for key in not_image_keys: for key in not_image_keys:
@ -512,6 +512,7 @@ def record_dataset(
stats = compute_stats(lerobot_dataset) stats = compute_stats(lerobot_dataset)
lerobot_dataset.stats = stats lerobot_dataset.stats = stats
else: else:
stats = {}
logging.info("Skipping computation of the dataset statistrics") logging.info("Skipping computation of the dataset statistrics")
hf_dataset = hf_dataset.with_format(None) # to remove transforms that cant be saved hf_dataset = hf_dataset.with_format(None) # to remove transforms that cant be saved