torch_em.util.util

  1import os
  2import warnings
  3from collections import OrderedDict
  4from typing import Optional, Tuple, Union
  5
  6import numpy as np
  7import torch
  8import torch_em
  9from matplotlib import colors
 10from numpy.typing import ArrayLike
 11
 12try:
 13    from torch._dynamo.eval_frame import OptimizedModule
 14except ImportError:
 15    OptimizedModule = None
 16
 17# torch doesn't support most unsigned types,
 18# so we map them to their signed equivalent
 19DTYPE_MAP = {
 20    np.dtype("uint16"): np.int16,
 21    np.dtype("uint32"): np.int32,
 22    np.dtype("uint64"): np.int64
 23}
 24"""@private
 25"""
 26
 27
 28def get_cache_directory() -> str:
 29    """Get the torch-em cache directory, for example for downloaded model weights.
 30
 31    Set the TORCH_EM_CACHEDIR environment variable to use a custom cache directory.
 32
 33    Returns:
 34        The path to the cache directory.
 35    """
 36    import pooch  # We import pooch here, so that 'import torch_em' does not need it.
 37
 38    default_cache_directory = os.path.expanduser(pooch.os_cache("torch_em"))
 39    return os.environ.get("TORCH_EM_CACHEDIR", default_cache_directory)
 40
 41
 42# This is a fairly brittle way to check if a module is compiled.
 43# Would be good to find a better solution, ideall something like model.is_compiled().
 44def is_compiled(model):
 45    """@private
 46    """
 47    if OptimizedModule is None:
 48        return False
 49    return isinstance(model, OptimizedModule)
 50
 51
 52def auto_compile(
 53    model: torch.nn.Module, compile_model: Optional[Union[str, bool]] = None, default_compile: bool = True
 54) -> torch.nn.Module:
 55    """Automatically compile a model for pytorch >= 2.
 56
 57    Args:
 58        model: The model.
 59        compile_model: Whether to comile the model.
 60            If None, it will not be compiled for torch < 2, and for torch > 2 the behavior
 61            specificed by 'default_compile' will be used. If a string is given it will be
 62            intepreted as the compile mode (torch.compile(model, mode=compile_model))
 63        default_compile: Whether to use the default compilation behavior for torch 2.
 64
 65    Returns:
 66        The compiled model.
 67    """
 68    torch_major = int(torch.__version__.split(".")[0])
 69
 70    if compile_model is None:
 71        if torch_major < 2:
 72            compile_model = False
 73        elif is_compiled(model):  # model is already compiled
 74            compile_model = False
 75        else:
 76            compile_model = default_compile
 77
 78    if compile_model:
 79        if torch_major < 2:
 80            raise RuntimeError("Model compilation is only supported for pytorch 2")
 81
 82        print("Compiling pytorch model ...")
 83        if isinstance(compile_model, str):
 84            model = torch.compile(model, mode=compile_model)
 85        else:
 86            model = torch.compile(model)
 87
 88    return model
 89
 90
 91def ensure_tensor(tensor: Union[torch.Tensor, ArrayLike], dtype: Optional[str] = None) -> torch.Tensor:
 92    """Ensure that the input is a torch tensor, by converting it if necessary.
 93
 94    Args:
 95        tensor: The input object, either a torch tensor or a numpy-array like object.
 96        dtype: The required data type for the output tensor.
 97
 98    Returns:
 99        The input, converted to a torch tensor if necessary.
100    """
101    if isinstance(tensor, np.ndarray):
102        if np.dtype(tensor.dtype) in DTYPE_MAP:
103            tensor = tensor.astype(DTYPE_MAP[tensor.dtype])
104        # Try to convert the tensor, even if it has wrong byte-order
105        try:
106            tensor = torch.from_numpy(tensor if tensor.flags.writeable else tensor.copy())
107        except ValueError:
108            tensor = tensor.view(tensor.dtype.newbyteorder())
109            if np.dtype(tensor.dtype) in DTYPE_MAP:
110                tensor = tensor.astype(DTYPE_MAP[tensor.dtype])
111            tensor = torch.from_numpy(tensor)
112
113    assert torch.is_tensor(tensor), f"Cannot convert {type(tensor)} to torch"
114    if dtype is not None:
115        tensor = tensor.to(dtype=dtype)
116    return tensor
117
118
119def validate_roi(roi, shape, patch_shape=None):
120    """Normalize an ROI to explicit slices and validate that it is non-empty."""
121    if roi is None:
122        return None
123    if isinstance(roi, slice):
124        roi = (roi,)
125    if not isinstance(roi, tuple):
126        raise TypeError(f"Invalid roi type: {type(roi)}")
127    if len(roi) > len(shape):
128        raise ValueError(f"Invalid roi {roi} for data shape {shape}: too many dimensions")
129
130    normalized_roi = []
131    for this_roi, dim in zip(roi, shape):
132        if not isinstance(this_roi, slice):
133            raise TypeError(f"Invalid roi entry: {this_roi}. Only slices are supported")
134        step = 1 if this_roi.step is None else this_roi.step
135        if step != 1:
136            raise ValueError(f"Invalid roi {roi}: slice steps other than 1 are not supported")
137        start, stop, _ = this_roi.indices(dim)
138        normalized_roi.append(slice(start, stop))
139
140    if len(roi) < len(shape):
141        normalized_roi.extend(slice(0, dim) for dim in shape[len(roi):])
142
143    roi_shape = tuple(sl.stop - sl.start for sl in normalized_roi)
144    if any(sh <= 0 for sh in roi_shape):
145        msg = f"Invalid roi {roi} for data shape {shape}: it results in an empty region"
146        if patch_shape is not None:
147            msg += f" for patch_shape {patch_shape}"
148        raise ValueError(msg)
149
150    return tuple(normalized_roi)
151
152
153def ensure_tensor_with_channels(
154    tensor: Union[torch.Tensor, ArrayLike], ndim: int, dtype: Optional[str] = None
155) -> torch.Tensor:
156    """Ensure that the input is a torch tensor of a given dimensionality with channels.
157
158    Args:
159        tensor: The input tensor or numpy-array like data.
160        ndim: The dimensionality of the output tensor.
161        dtype: The data type of the output tensor.
162
163    Returns:
164        The input converted to a torch tensor of the requested dimensionality.
165    """
166    assert ndim in (2, 3, 4), f"{ndim}"
167    tensor = ensure_tensor(tensor, dtype)
168    if ndim == 2:
169        assert tensor.ndim in (2, 3, 4, 5), f"{tensor.ndim}"
170        if tensor.ndim == 2:
171            tensor = tensor[None]
172        elif tensor.ndim == 4:
173            assert tensor.shape[0] == 1, f"{tensor.shape}"
174            tensor = tensor[0]
175        elif tensor.ndim == 5:
176            assert tensor.shape[:2] == (1, 1), f"{tensor.shape}"
177            tensor = tensor[0, 0]
178    elif ndim == 3:
179        assert tensor.ndim in (3, 4, 5), f"{tensor.ndim}"
180        if tensor.ndim == 3:
181            tensor = tensor[None]
182        elif tensor.ndim == 5:
183            assert tensor.shape[0] == 1, f"{tensor.shape}"
184            tensor = tensor[0]
185    else:
186        assert tensor.ndim in (4, 5), f"{tensor.ndim}"
187        if tensor.ndim == 5:
188            assert tensor.shape[0] == 1, f"{tensor.shape}"
189            tensor = tensor[0]
190    return tensor
191
192
193def ensure_array(array: Union[np.ndarray, torch.Tensor], dtype: str = None) -> np.ndarray:
194    """Ensure that the input is a numpy array, by converting it if necessary.
195
196    Args:
197        array: The input torch tensor or numpy array.
198        dtype: The dtype of the ouptut array.
199
200    Returns:
201        The input converted to a numpy array if necessary.
202    """
203    if torch.is_tensor(array):
204        array = array.detach().cpu().numpy()
205    assert isinstance(array, np.ndarray), f"Cannot convert {type(array)} to numpy"
206    if dtype is not None:
207        array = np.require(array, dtype=dtype)
208    return array
209
210
211def ensure_spatial_array(array: Union[np.ndarray, torch.Tensor], ndim: int, dtype: str = None) -> np.ndarray:
212    """Ensure that the input is a numpy array of a given dimensionality.
213
214    Args:
215        array: The input numpy array or torch tensor.
216        ndim: The requested dimensionality.
217        dtype: The dtype of the output array.
218
219    Returns:
220        A numpy array of the requested dimensionality and data type.
221    """
222    assert ndim in (2, 3)
223    array = ensure_array(array, dtype)
224    if ndim == 2:
225        assert array.ndim in (2, 3, 4, 5), str(array.ndim)
226        if array.ndim == 3:
227            assert array.shape[0] == 1
228            array = array[0]
229        elif array.ndim == 4:
230            assert array.shape[:2] == (1, 1)
231            array = array[0, 0]
232        elif array.ndim == 5:
233            assert array.shape[:3] == (1, 1, 1)
234            array = array[0, 0, 0]
235    else:
236        assert array.ndim in (3, 4, 5), str(array.ndim)
237        if array.ndim == 4:
238            assert array.shape[0] == 1, f"{array.shape}"
239            array = array[0]
240        elif array.ndim == 5:
241            assert array.shape[:2] == (1, 1)
242            array = array[0, 0]
243    return array
244
245
246def ensure_patch_shape(
247    raw: np.ndarray,
248    labels: Optional[np.ndarray],
249    patch_shape: Tuple[int, ...],
250    have_raw_channels: bool = False,
251    have_label_channels: bool = False,
252    channel_first: bool = True,
253) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray]:
254    """Ensure that the raw data and labels have at least the requested patch shape.
255
256    If either raw data or labels do not have the patch shape they will be padded.
257
258    Args:
259        raw: The input raw data.
260        labels: The input labels.
261        patch_shape: The required minimal patch shape.
262        have_raw_channels: Whether the raw data has channels.
263        have_label_channels: Whether the label data has channels.
264        channel_first: Whether the channel axis is the first or last axis.
265
266    Returns:
267        The raw data.
268        The labels.
269    """
270    raw_shape = raw.shape
271    if labels is not None:
272        labels_shape = labels.shape
273
274    # In case the inputs has channels and they are channels first
275    # IMPORTANT: for ImageCollectionDataset
276    if have_raw_channels and channel_first:
277        raw_shape = raw_shape[1:]
278
279    if have_label_channels and channel_first and labels is not None:
280        labels_shape = labels_shape[1:]
281
282    # Extract the pad_width and pad the raw inputs
283    if any(sh < psh for sh, psh in zip(raw_shape, patch_shape)):
284        pw = [(0, max(0, psh - sh)) for sh, psh in zip(raw_shape, patch_shape)]
285
286        if have_raw_channels and channel_first:
287            pad_width = [(0, 0), *pw]
288        elif have_raw_channels and not channel_first:
289            pad_width = [*pw, (0, 0)]
290        else:
291            pad_width = pw
292
293        raw = np.pad(array=raw, pad_width=pad_width)
294
295    # Extract the pad width and pad the label inputs
296    if labels is not None and any(sh < psh for sh, psh in zip(labels_shape, patch_shape)):
297        pw = [(0, max(0, psh - sh)) for sh, psh in zip(labels_shape, patch_shape)]
298
299        if have_label_channels and channel_first:
300            pad_width = [(0, 0), *pw]
301        elif have_label_channels and not channel_first:
302            pad_width = [*pw, (0, 0)]
303        else:
304            pad_width = pw
305
306        labels = np.pad(array=labels, pad_width=pad_width)
307    if labels is None:
308        return raw
309    else:
310        return raw, labels
311
312
313def get_constructor_arguments(obj):
314    """@private
315    """
316    # All relevant torch_em classes have 'init_kwargs' to directly recover the init call.
317    if hasattr(obj, "init_kwargs"):
318        return getattr(obj, "init_kwargs")
319
320    def _get_args(obj, param_names):
321        return {name: getattr(obj, name) for name in param_names}
322
323    # We don't need to find the constructor arguments for optimizers/schedulers because we deserialize the state later.
324    if isinstance(
325        obj, (
326            torch.optim.Optimizer,
327            torch.optim.lr_scheduler._LRScheduler,
328            # ReduceLROnPlateau does not inherit from _LRScheduler
329            torch.optim.lr_scheduler.ReduceLROnPlateau
330        )
331    ):
332        return {}
333
334    # recover the arguments for torch dataloader
335    elif isinstance(obj, torch.utils.data.DataLoader):
336        # These are all the "simple" arguements.
337        # "sampler", "batch_sampler" and "worker_init_fn" are more complicated
338        # and generally not used in torch_em
339        sampler = getattr(obj, "sampler", None)
340        if sampler is not None and not isinstance(
341            sampler,
342            (
343                torch.utils.data.RandomSampler,
344                torch.utils.data.SequentialSampler,
345                torch.utils.data.SubsetRandomSampler,
346            ),
347        ):
348            warnings.warn(
349                f"DataLoader uses sampler {type(sampler).__name__}, but only its effective `shuffle` setting "
350                "is serialized. `DefaultTrainer.from_checkpoint` will recreate the loader without the original "
351                "sampler, so sampling behavior may change."
352            )
353        shuffle = getattr(obj, "shuffle", None)
354        if shuffle is None:
355            shuffle = getattr(sampler, "shuffle", None)
356        if shuffle is None:
357            # Only randomized samplers map to shuffle=True. SequentialSampler is handled
358            # by the default fallback of shuffle=False and does not need a special case.
359            shuffle = isinstance(sampler, (torch.utils.data.RandomSampler, torch.utils.data.SubsetRandomSampler))
360
361        return {
362            **_get_args(
363                obj, [
364                    "batch_size", "num_workers", "pin_memory", "drop_last",
365                    "persistent_workers", "prefetch_factor", "timeout"
366                ]
367            ),
368            "shuffle": shuffle,
369        }
370
371    # TODO support common torch losses (e.g. CrossEntropy, BCE)
372    warnings.warn(
373        f"Constructor arguments for {type(obj)} cannot be deduced.\n" +
374        "For this object, empty constructor arguments will be used.\n" +
375        "The trainer can probably not be correctly deserialized via 'DefaultTrainer.from_checkpoint'."
376    )
377    return {}
378
379
380def get_trainer(checkpoint: str, name: str = "best", device: Optional[str] = None):
381    """Load trainer from a checkpoint.
382
383    Args:
384        checkpoint: The path to the checkpoint.
385        name: The name of the checkpoint.
386        device: The device to use for loading the checkpoint.
387
388    Returns:
389        The trainer.
390    """
391    # try to load from file
392    if isinstance(checkpoint, str):
393        assert os.path.exists(checkpoint), checkpoint
394        trainer = torch_em.trainer.DefaultTrainer.from_checkpoint(checkpoint, name=name, device=device)
395    else:
396        trainer = checkpoint
397    assert isinstance(trainer, torch_em.trainer.DefaultTrainer)
398    return trainer
399
400
401def get_normalizer(trainer):
402    """@private
403    """
404    dataset = trainer.train_loader.dataset
405    while (
406        isinstance(dataset, torch_em.data.concat_dataset.ConcatDataset) or
407        isinstance(dataset, torch.utils.data.dataset.ConcatDataset)
408    ):
409        dataset = dataset.datasets[0]
410
411    if isinstance(dataset, torch.utils.data.dataset.Subset):
412        dataset = dataset.dataset
413
414    preprocessor = dataset.raw_transform
415
416    if hasattr(preprocessor, "normalizer"):
417        return preprocessor.normalizer
418    else:
419        return preprocessor
420
421
422def load_model(
423    checkpoint: str,
424    model: Optional[torch.nn.Module] = None,
425    name: str = "best",
426    state_key: Optional[str] = "model_state",
427    device: Optional[str] = None,
428) -> torch.nn.Module:
429    """Load a model from a trainer checkpoint or a serialized torch model.
430
431    This function can either load the model directly (`model` is not passed),
432    or deserialize the model state and then load it (`model` is passed).
433
434    The `checkpoint` argument must either point to the checkpoint directory of a torch_em trainer
435    or to a serialized torch model.
436
437    Args:
438        checkpoint: The path to the checkpoint folder or serialized torch model.
439        model: The model for which the state should be loaded.
440            If it is not passed, the model class and parameters will also be loaded from the trainer.
441        name: The name of the checkpoint.
442        state_key: The name of the model state to load. Set to None if the model state is stored top-level.
443        device: The device on which to load the model.
444
445    Returns:
446        The model.
447    """
448    if model is None and os.path.isdir(checkpoint):  # Load the model and its state from a torch_em checkpoint.
449        model = get_trainer(checkpoint, name=name, device=device).model
450
451    elif model is None:  # Load the model from a serialized model.
452        model = torch.load(checkpoint, map_location=device, weights_only=False)
453
454    else:  # Load the model state from a checkpoint.
455        if os.path.isdir(checkpoint):  # From a torch_em checkpoint.
456            ckpt = os.path.join(checkpoint, f"{name}.pt")
457        else:  # From a serialized path.
458            ckpt = checkpoint
459
460        state = torch.load(ckpt, map_location=device, weights_only=False)
461        if state_key is not None:
462            state = state[state_key]
463
464        # To enable loading compiled models.
465        compiled_prefix = "_orig_mod."
466        state = OrderedDict(
467            [(k[len(compiled_prefix):] if k.startswith(compiled_prefix) else k, v) for k, v in state.items()]
468        )
469
470        model.load_state_dict(state)
471        if device is not None:
472            model.to(device)
473
474    return model
475
476
477def model_is_equal(model1, model2):
478    """@private
479    """
480    for p1, p2 in zip(model1.parameters(), model2.parameters()):
481        if p1.data.ne(p2.data).sum() > 0:
482            return False
483    return True
484
485
486def get_random_colors(labels: np.ndarray) -> colors.ListedColormap:
487    """Generate a random color map for a label image.
488
489    Args:
490        labels: The labels.
491
492    Returns:
493        The color map.
494    """
495    unique_labels = np.unique(labels)
496    have_zero = 0 in unique_labels
497    cmap = [[0, 0, 0]] if have_zero else []
498    cmap += np.random.rand(len(unique_labels), 3).tolist()
499    cmap = colors.ListedColormap(cmap)
500    return cmap