torch_em.data.pseudo_label_dataset

 1import os
 2from typing import Union, Tuple, Optional, List, Any, Callable
 3
 4import torch
 5
 6from .raw_dataset import RawDataset
 7from ..util import ensure_tensor_with_channels
 8
 9
10class PseudoLabelDataset(RawDataset):
11    """Dataset that uses a prediction function to provide raw data and pseudo labels for segmentation training.
12
13    The dataset loads a patch from the raw data and then applies the pseudo labeler to it to predict pseudo labels.
14    The raw data and pseudo labels are returned together as a sample for a batch.
15    The datataset supports all file formats that can be opened with `elf.io.open_file`, such as hdf5, zarr or n5.
16    Use `raw_path` to specify the path to the file and `raw_key` to specify the internal dataset.
17    It also supports regular image formats, such as .tif. For these cases set `raw_key=None`.
18
19    Args:
20        raw_path: The file path to the raw image data. May also be a list of file paths.
21        raw_key: The key to the internal dataset containing the raw data.
22        patch_shape: The patch shape for a training sample.
23        pseudo_labeler: The pseudo labeler. Must be a function that accepts the raw data as torch tensor
24            and that returns the predicted labels as torch tensor.
25        raw_transform: Transformation applied to the raw data of a sample.
26        label_transform: Transformation applied to the label data of a sample.
27        roi: Region of interest in the raw data.
28            If given, the raw data will only be loaded from the corresponding area.
29        dtype: The return data type of the raw data.
30        n_samples: The length of this dataset. If None, the length will be set to `len(raw_image_paths)`.
31        sampler: Sampler for rejecting samples according to a defined criterion.
32            The sampler must be a callable that accepts the raw data and label data (as numpy arrays) as input.
33        ndim: The spatial dimensionality of the data. If None, will be derived from the raw data.
34        with_channels: Whether the raw data has channels.
35        labeler_device: The expected device for the pseudo labeler.
36    """
37    def __init__(
38        self,
39        raw_path: Union[List[Any], str, os.PathLike],
40        raw_key: Optional[str],
41        patch_shape: Tuple[int, ...],
42        pseudo_labeler: Callable,
43        raw_transform: Optional[Callable] = None,
44        label_transform: Optional[Callable] = None,
45        transform: Optional[Callable] = None,
46        roi: Optional[Union[slice, Tuple[slice, ...]]] = None,
47        dtype: torch.dtype = torch.float32,
48        n_samples: Optional[int] = None,
49        sampler: Optional[Callable] = None,
50        ndim: Optional[Union[int]] = None,
51        with_channels: bool = False,
52        labeler_device: Optional[Union[str, torch.device]] = None,
53    ):
54        super().__init__(
55            raw_path, raw_key, patch_shape, raw_transform=raw_transform, transform=transform, roi=roi,
56            dtype=dtype, n_samples=n_samples, sampler=sampler, ndim=ndim, with_channels=with_channels
57        )
58        self.pseudo_labeler = pseudo_labeler
59        self.label_transform = label_transform
60        self.labeler_device = next(pseudo_labeler.parameters()).device if labeler_device is None else labeler_device
61
62    def __getitem__(self, index):
63        raw = self._get_sample(index)
64
65        # Transform for augmentations.
66        # Applied to the raw data since, labels are generated on the fly by the pseudo_labeler.
67        if self.transform is not None:
68            raw = self.transform(raw)[0]
69            if self.trafo_halo is not None:
70                raw = self.crop(raw)
71
72        raw = ensure_tensor_with_channels(raw, ndim=self._ndim, dtype=self.dtype)
73        with torch.no_grad():
74            # Ilastik needs uint as input, so normalize afterwards.
75            labels = self.pseudo_labeler(raw[None].to(self.labeler_device))[0]
76
77        # Normalize the raw data.
78        if self.raw_transform is not None:
79            raw = self.raw_transform(raw.cpu().detach().numpy())
80        raw = ensure_tensor_with_channels(raw, ndim=self._ndim, dtype=self.dtype)
81
82        if self.label_transform is not None:
83            labels = self.label_transform(labels)
84        labels = ensure_tensor_with_channels(labels, ndim=self._ndim)
85
86        return raw, labels