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

Fix paths & add add_frame doc

Browse Source
This commit is contained in:
Simon Alibert
2024-10-22 22:46:34 +02:00
parent 6c2cb6e107
commit 237a484be0
2 changed files with 22 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
lerobot/common/datasets/image_writer.py
View File

@@ -99,19 +99,16 @@ class ImageWriter:
img = Image.fromarray(image.numpy())
img.save(str(file_path), quality=100)
def get_image_file_path(
self, episode_index: int, image_key: str, frame_index: int, return_str: bool = True
) -> str | Path:
def get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
fpath = self.image_path.format(
image_key=image_key, episode_index=episode_index, frame_index=frame_index
)
return str(self.dir / fpath) if return_str else self.dir / fpath
return self.dir / fpath
def get_episode_dir(self, episode_index: int, image_key: str, return_str: bool = True) -> str | Path:
dir_path = self.get_image_file_path(
episode_index=episode_index, image_key=image_key, frame_index=0, return_str=False
def get_episode_dir(self, episode_index: int, image_key: str) -> Path:
return self.get_image_file_path(
episode_index=episode_index, image_key=image_key, frame_index=0
).parent
return str(dir_path) if return_str else dir_path
def stop(self, timeout=20) -> None:
"""Stop the image writer, waiting for all processes or threads to finish."""