Add test_image_writer, accept PIL images, improve ImageWriter perf in main process

2024-11-02 20:00:07 +01:00
parent 375abd3020
commit 6b2ec1ed77
4 changed files with 426 additions and 16 deletions
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@@ -590,7 +590,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
                img_path.parent.mkdir(parents=True, exist_ok=True)

            self.image_writer.save_image(
-                image_array=frame[cam_key],
+                image=frame[cam_key],
                fpath=img_path,
            )