@@ -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