"""
High-level file handlers (load/save) for generic and hierarchical file formats
"""
import logging
import os
from typing import List, Optional, Tuple
import numpy as np
from .generic import load_npy, load_raw, load_tif, write_npy, write_tif
from .hierarchy import load_ND_with_key, write_hdf5
logger = logging.getLogger(__name__)
[docs]
def load_array_from_file(
path: str,
key: Optional[str] = None,
ndims: Optional[List[int]] = None,
shape: Optional[Tuple[int, ...]] = None,
offset: Optional[int] = None,
gap: Optional[int] = None,
) -> np.ndarray:
"""
Load array from a file. The file type is inferred from the extension.
Currently supports .tif, .tiff, .npy, .mat, .h5, .hdf5, and .raw.
Args:
path (str): Path to the file.
key (str): Key to specify the dataset (optional).
ndims (list): List of desired dimensions for filtering datasets.
shape (tuple): Shape of the data for .raw files (optional).
offset (int): Offset for .raw files (optional).
gap (int): Gap for .raw files (optional).
Returns:
numpy.ndarray: The loaded array.
Raises:
ValueError: If the file type is unsupported or no valid dataset is found.
"""
file_path = path # The function signature is simplified for users, although I think file_path is clearer
# Check file existence
if not os.path.exists(file_path):
raise FileNotFoundError(f"The specified file '{file_path}' does not exist. Please check your file path and working directory.")
# Infer file type from extension
_, ext = os.path.splitext(file_path)
ext = ext.lower()
if ext in [".tif", ".tiff"]:
return load_tif(file_path)
elif ext == ".npy":
return load_npy(file_path)
elif ext in [".mat", ".h5", ".hdf5"]:
return load_ND_with_key(file_path, key, ndims)
elif ext == ".raw":
if shape is None:
raise ValueError(
f"Please at least provide 'shape' of the expected data array to correctly load the .raw file {file_path}."
)
raw_args = {"shape": shape, "offset": offset, "gap": gap}
raw_args = {
k: v for k, v in raw_args.items() if v is not None
} # Remove argument with None
return load_raw(file_path, **raw_args)
else:
raise ValueError(
f"Unsupported file type: '{ext}'. Supported types are .tif, .tiff, .mat, .h5, .hdf5, .npy, and .raw."
)
[docs]
def save_array(data, file_dir='', file_name='ptyrad_init_meas', file_format="hdf5", output_shape=None, append_shape=True, **kwargs):
"""
Save an ND array to the specified file format.
Args:
data (numpy.ndarray): ND array to save.
file_dir (str): Directory to save the file.
file_name (str): Base name of the file (without extension).
file_format (str): File format to save as ("tif", "npy", "hdf5", "mat").
output_shape (tuple, optional): Desired shape for the output array.
append_shape (bool): Whether to append the array shape to the filename.
**kwargs: Additional arguments for specific file formats.
"""
# Reshape data if output_ndim is specified
if output_shape is not None:
try:
data = data.reshape(output_shape)
except ValueError as e:
logger.warning(f"WARNING: {e}, the data shape is preserved as {data.shape}")
# Append shape to the filename if enabled
if append_shape:
shape_str = "_".join(map(str, data.shape))
file_name = f"{file_name}_{shape_str}"
# Construct the full file path
file_format = file_format.lower()
file_path = os.path.join(file_dir, f"{file_name}.{file_format}")
logger.info(f"Saving array with shape = {data.shape} and dtype = {data.dtype}")
if os.path.isfile(file_path):
logger.info(f"file path = '{file_path}' already exists, the file will be overwritten.")
if file_format in ["tif", "tiff"]:
write_tif(file_path, data)
elif file_format == "npy":
write_npy(file_path, data)
elif file_format in ["hdf5", "h5", "mat"]:
# Saving .mat into hdf5 as if it were .mat v7.3. This ensures compatibility with py4DGUI.
write_hdf5(file_path, data, **kwargs)
else:
raise ValueError(f"Unsupported file format: {file_format}")