Source code for brainlit.registration.file_io

import numpy as np
import SimpleITK as sitk
from pathlib import Path
import pickle

def _validate_inputs(**kwargs):
    """Accepts arbitrary kwargs. If recognized, they are validated.
    If they cannot be validated an exception is raised.
    A dictionary containing the validated kwargs is returned."""

    # Validate data.

    if "data" in kwargs.keys():
        data = kwargs["data"]
        if isinstance(data, dict):
            # Verify that each value in the dict is a np.ndarray object and recursion will stop at 1 level.
            if not all(map(lambda datum: isinstance(datum, np.ndarray), data.values())):
                raise TypeError(
                    f"data must be either a np.ndarray or a dictionary mapping only to np.ndarray objects."
            # Recurse into each element of the dictionary keys of np.ndarray objects.
            for key, datum in data.items():
                data[key] = _validate_inputs(data=datum)["data"]
        # If data is neither a list, dict, nor a np.ndarray, raise TypeError.
        elif not isinstance(data, np.ndarray):
            raise TypeError(f"data must be a np.ndarray.\ntype(data): {type(data)}.")
        # data is a np.ndarray.

    # Validate file_path.

    if "file_path" in kwargs.keys():
        file_path = kwargs["file_path"]
        file_path = (
        )  # Will raise exceptions if file_path is not an oppropriate type.
        if not file_path.parent.is_dir():
            raise FileNotFoundError(
                f"file_path corresponds to a location that does not presently exist.\n"
                f"file_path.parent: {file_path.parent}."
        # Validate extension.
        if not file_path.suffix:
            default_extension = ".vtk"
            file_path = file_path.with_suffix(default_extension)
        # file_path is valid.

    return kwargs

[docs]def save(data, file_path): """ Save data to file_path Args: data (np.ndarray, dict): An array or a dictionary with np.ndarray values to be saved. file_path (str, Path): The file path that data will be saved to. Accepts an arbitrary suffix but provides .vtk by default. Raises: Exception: Raised if _validate_inputs has failed to catch an improper type for data. """ # Validate inputs. inputs = {"data": data, "file_path": file_path} validated_inputs = _validate_inputs(**inputs) data = validated_inputs["data"] file_path = validated_inputs["file_path"] # data is either a single np.ndarray or a list of np.ndarrays. if isinstance(data, np.ndarray): # Convert data to sitk.Image. # Take the transpose of data to premptively correct for the f-ordering of SimpleITK Images. data = data.T data_Image = sitk.GetImageFromArray(data) # Side effect: breaks alias. # Save data to file_path. sitk.WriteImage(data_Image, str(file_path)) # If data is a dictionary it must map to np.ndarray objects. elif isinstance(data, dict): np.savez(file_path.with_suffix(""), **data) # '.npz' is appended. else: # _validate_inputs has failed. raise Exception( f"_validate_inputs has failed to prevent an improper type for data.\n" f"type(data): {type(data)}." )
[docs]def load(file_path): """ Load data from file_path. Args: file_path (str, Path): The file path from which data will be retrieved. Returns: np.ndarray, dict: The array or dict of arrays saved at file_path. """ # Validate inputs. inputs = {"file_path": file_path} validated_inputs = _validate_inputs(**inputs) file_path = validated_inputs["file_path"] if file_path.suffix == ".npz": data = np.load(file_path) # data is a dictionary. return data else: # Read in data as sitk.Image. data_Image = sitk.ReadImage(str(file_path)) # Convert data_Image to np.ndarray. data = sitk.GetArrayFromImage(data_Image) # Take the transpose of data to correct for the f-ordering of SimpleITK Images. data = data.T # data is a np.ndarray. return data
def save_pickled(obj, file_path): """ Pickle object obj and save it to file_path. Args: obj (object): The pickleable object to be saved. file_path (str, Path): The file path at which to save obj. """ # Validate file_path. file_path = _validate_inputs(file_path=file_path)["file_path"] with open(file_path, "wb") as file: pickle.dump(obj, file) def load_pickled(file_path): """ Load pickled object from file_path. Args: file_path (str, Path): The file path at which a pickled object is saved. Returns: object: The pickled object saved at file_path. """ # Validate file_path. file_path = _validate_inputs(file_path=file_path)["file_path"] # TODO: verify behavior of suffixes. with open(file_path, "rb") as file: return pickle.load(file)