Skip to content

Commit 9a49e57

Browse files
refactor(datasets): add compress_level parameter to write_image() and set it to 1 (huggingface#2135)
* refactor(datasets): add compress_level parameter to write_image() and set it to 1 * docs(dataset): add docs to write_image()
1 parent 6c28ef8 commit 9a49e57

1 file changed

Lines changed: 25 additions & 2 deletions

File tree

src/lerobot/datasets/image_writer.py

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,38 @@ def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True)
6868
return PIL.Image.fromarray(image_array)
6969

7070

71-
def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path):
71+
def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1):
72+
"""
73+
Saves a NumPy array or PIL Image to a file.
74+
75+
This function handles both NumPy arrays and PIL Image objects, converting
76+
the former to a PIL Image before saving. It includes error handling for
77+
the save operation.
78+
79+
Args:
80+
image (np.ndarray | PIL.Image.Image): The image data to save.
81+
fpath (Path): The destination file path for the image.
82+
compress_level (int, optional): The compression level for the saved
83+
image, as used by PIL.Image.save(). Defaults to 1.
84+
Refer to: https://github.com/huggingface/lerobot/pull/2135
85+
for more details on the default value rationale.
86+
87+
Raises:
88+
TypeError: If the input 'image' is not a NumPy array or a
89+
PIL.Image.Image object.
90+
91+
Side Effects:
92+
Prints an error message to the console if the image writing process
93+
fails for any reason.
94+
"""
7295
try:
7396
if isinstance(image, np.ndarray):
7497
img = image_array_to_pil_image(image)
7598
elif isinstance(image, PIL.Image.Image):
7699
img = image
77100
else:
78101
raise TypeError(f"Unsupported image type: {type(image)}")
79-
img.save(fpath)
102+
img.save(fpath, compress_level=compress_level)
80103
except Exception as e:
81104
print(f"Error writing image {fpath}: {e}")
82105

0 commit comments

Comments
 (0)