tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix visualize_dataset

Browse Source
This commit is contained in:
Cadene
2024-04-19 23:40:35 +00:00
parent 9b4c2e2a9f
commit 85f1554da8
1 changed files with 15 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
lerobot/scripts/visualize_dataset.py
View File

@@ -22,11 +22,22 @@ def visualize_dataset_cli(cfg: dict):
def cat_and_write_video(video_path, frames, fps): def cat_and_write_video(video_path, frames, fps):
# Expects images in [0, 255].
frames = torch.cat(frames) frames = torch.cat(frames)
assert frames.dtype == torch.uint8
frames = einops.rearrange(frames, "b c h w -> b h w c").numpy() # Expects images in [0, 1].
imageio.mimsave(video_path, frames, fps=fps) frame = frames[0]
_, c, h, w = frame.shape
assert c < h and c < w, f"expect channel first images, but instead {frame.shape}"
# sanity check that images are float32 in range [0,1]
assert frame.dtype == torch.float32, f"expect torch.float32, but instead {frame.dtype=}"
assert frame.max() <= 1, f"expect pixels lower than 1, but instead {frame.max()=}"
assert frame.min() >= 0, f"expect pixels greater than 1, but instead {frame.min()=}"
# convert to channel last uint8 [0, 255]
frames = einops.rearrange(frames, "b c h w -> b h w c")
frames = (frames * 255).type(torch.uint8)
imageio.mimsave(video_path, frames.numpy(), fps=fps)
def visualize_dataset(cfg: dict, out_dir=None): def visualize_dataset(cfg: dict, out_dir=None):