torch_em.util.image

  1import os
  2from typing import Optional, Sequence, Union
  3
  4import imageio.v3 as imageio
  5import numpy as np
  6from elf.io import open_file
  7from numpy.typing import ArrayLike
  8
  9try:
 10    import tifffile
 11except ImportError:
 12    tifffile = None
 13
 14TIF_EXTS = (".tif", ".tiff")
 15
 16
 17def supports_memmap(image_path):
 18    """@private
 19    """
 20    if tifffile is None:
 21        return False
 22    ext = os.path.splitext(image_path)[1]
 23    if ext.lower() not in TIF_EXTS:
 24        return False
 25    try:
 26        tifffile.memmap(image_path, mode="r")
 27    except ValueError:
 28        return False
 29    return True
 30
 31
 32def load_image(image_path, memmap=True):
 33    """@private
 34    """
 35    if supports_memmap(image_path) and memmap:
 36        return tifffile.memmap(image_path, mode="r")
 37    elif tifffile is not None and os.path.splitext(image_path)[1].lower() in (".tiff", ".tif"):
 38        return tifffile.imread(image_path)
 39    elif os.path.splitext(image_path)[1].lower() == ".nrrd":
 40        import nrrd
 41        return nrrd.read(image_path)[0]
 42    elif os.path.splitext(image_path)[1].lower() == ".mha":
 43        import SimpleITK as sitk
 44        image = sitk.ReadImage(image_path)
 45        return sitk.GetArrayFromImage(image)
 46    else:
 47        return imageio.imread(image_path)
 48
 49
 50class MultiDatasetWrapper:
 51    """@private
 52    """
 53    def __init__(self, *file_datasets):
 54        # Make sure we have the same shapes.
 55        reference_shape = file_datasets[0].shape
 56        assert all(reference_shape == ds.shape for ds in file_datasets)
 57        self.file_datasets = file_datasets
 58
 59        self.shape = (len(self.file_datasets),) + reference_shape
 60
 61    def __getitem__(self, index):
 62        channel_index, spatial_index = index[:1], index[1:]
 63        data = []
 64        for ds in self.file_datasets:
 65            ds_data = ds[spatial_index]
 66            data.append(ds_data)
 67        data = np.stack(data)
 68        data = data[channel_index]
 69        return data
 70
 71
 72def load_data(
 73    path: Union[str, Sequence[str]],
 74    key: Optional[Union[str, Sequence[str]]] = None,
 75    mode: str = "r",
 76) -> ArrayLike:
 77    """Load data from a file or multiple files.
 78
 79    Supports loading regular image formats, such as tif or jpg, or container data formats, such as hdf5, n5 or zarr.
 80    For the latter case, specify the name of the internal dataset to load via the `key` argument.
 81
 82    Args:
 83        path: The file path or paths to the data.
 84        key: The key or keys to the internal datasets.
 85        mode: The mode for reading datasets.
 86
 87    Returns:
 88        The loaded data.
 89    """
 90    have_single_file = isinstance(path, str)
 91    have_single_key = isinstance(key, str)
 92
 93    if key is None:
 94        if have_single_file:
 95            return load_image(path)
 96        else:
 97            return np.stack([load_image(p) for p in path])
 98    else:
 99        if have_single_key and have_single_file:
100            return open_file(path, mode=mode)[key]
101        elif have_single_key and not have_single_file:
102            return MultiDatasetWrapper(*[open_file(p, mode=mode)[key] for p in path])
103        elif not have_single_key and have_single_file:
104            return MultiDatasetWrapper(*[open_file(path, mode=mode)[k] for k in key])
105        else:  # have multipe keys and multiple files
106            return MultiDatasetWrapper(*[open_file(p, mode=mode)[k] for k in key for p in path])