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    # mrc files require key="data"; set it automatically if no key was provided
 94    if key is None:
 95        if have_single_file and str(path).endswith(".mrc"):
 96            key = "data"
 97            have_single_key = True
 98        elif not have_single_file and all(str(p).endswith(".mrc") for p in path):
 99            key = "data"
100            have_single_key = True
101
102    if key is None:
103        if have_single_file:
104            return load_image(path)
105        else:
106            return np.stack([load_image(p) for p in path])
107    else:
108        if have_single_key and have_single_file:
109            return open_file(path, mode=mode)[key]
110        elif have_single_key and not have_single_file:
111            return MultiDatasetWrapper(*[open_file(p, mode=mode)[key] for p in path])
112        elif not have_single_key and have_single_file:
113            return MultiDatasetWrapper(*[open_file(path, mode=mode)[k] for k in key])
114        else:  # have multipe keys and multiple files
115            return MultiDatasetWrapper(*[open_file(p, mode=mode)[k] for k in key for p in path])