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])
TIF_EXTS =
('.tif', '.tiff')
def
load_data( path: Union[str, Sequence[str]], key: Union[str, Sequence[str], NoneType] = None, mode: str = 'r') -> Union[Buffer, numpy._typing._array_like._SupportsArray[numpy.dtype[Any]], numpy._typing._nested_sequence._NestedSequence[numpy._typing._array_like._SupportsArray[numpy.dtype[Any]]], bool, int, float, complex, str, bytes, numpy._typing._nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]]:
73def load_data( 74 path: Union[str, Sequence[str]], 75 key: Optional[Union[str, Sequence[str]]] = None, 76 mode: str = "r", 77) -> ArrayLike: 78 """Load data from a file or multiple files. 79 80 Supports loading regular image formats, such as tif or jpg, or container data formats, such as hdf5, n5 or zarr. 81 For the latter case, specify the name of the internal dataset to load via the `key` argument. 82 83 Args: 84 path: The file path or paths to the data. 85 key: The key or keys to the internal datasets. 86 mode: The mode for reading datasets. 87 88 Returns: 89 The loaded data. 90 """ 91 have_single_file = isinstance(path, str) 92 have_single_key = isinstance(key, str) 93 94 if key is None: 95 if have_single_file: 96 return load_image(path) 97 else: 98 return np.stack([load_image(p) for p in path]) 99 else: 100 if have_single_key and have_single_file: 101 return open_file(path, mode=mode)[key] 102 elif have_single_key and not have_single_file: 103 return MultiDatasetWrapper(*[open_file(p, mode=mode)[key] for p in path]) 104 elif not have_single_key and have_single_file: 105 return MultiDatasetWrapper(*[open_file(path, mode=mode)[k] for k in key]) 106 else: # have multipe keys and multiple files 107 return MultiDatasetWrapper(*[open_file(p, mode=mode)[k] for k in key for p in path])
Load data from a file or multiple files.
Supports loading regular image formats, such as tif or jpg, or container data formats, such as hdf5, n5 or zarr.
For the latter case, specify the name of the internal dataset to load via the key
argument.
Arguments:
- path: The file path or paths to the data.
- key: The key or keys to the internal datasets.
- mode: The mode for reading datasets.
Returns:
The loaded data.