synapse_net.training.dataloaders

  1from typing import Optional, Tuple, Union
  2
  3import torch
  4import torch_em
  5from torch_em.data import RawDatasetWithMasks
  6
  7
  8def _adjust_patch_shape(data_shape, patch_shape):
  9    # If data is 2D and patch_shape is 3D, drop the extra dimension in patch_shape
 10    if data_shape == 2 and len(patch_shape) == 3:
 11        return patch_shape[1:]  # Remove the leading dimension in patch_shape
 12    return patch_shape  # Return the original patch_shape for 3D data
 13
 14
 15def _determine_ndim(patch_shape):
 16    # Check for 2D or 3D training
 17    try:
 18        z, y, x = patch_shape
 19    except ValueError:
 20        y, x = patch_shape
 21        z = 1
 22    is_2d = z == 1
 23    ndim = 2 if is_2d else 3
 24    return is_2d, ndim
 25
 26
 27def get_supervised_loader(
 28    data_paths: Tuple[str],
 29    raw_key: str,
 30    label_key: str,
 31    patch_shape: Tuple[int, int, int],
 32    batch_size: int,
 33    n_samples: Optional[int],
 34    add_boundary_transform: bool = True,
 35    label_dtype=torch.float32,
 36    rois: Optional[Tuple[Tuple[slice]]] = None,
 37    sampler: Optional[Union[callable, bool]] = None,
 38    ignore_label: Optional[int] = None,
 39    label_transform: Optional[callable] = None,
 40    label_paths: Optional[Tuple[str]] = None,
 41    raw_transform: Optional[callable] = None,
 42    **loader_kwargs,
 43) -> torch.utils.data.DataLoader:
 44    """Get a dataloader for supervised segmentation training.
 45
 46    Args:
 47        data_paths: The filepaths to the hdf5 files containing the training data.
 48        raw_key: The key that holds the raw data inside of the hdf5.
 49        label_key: The key that holds the labels inside of the hdf5.
 50        patch_shape: The patch shape used for a training example.
 51            In order to run 2d training pass a patch shape with a singleton in the z-axis,
 52            e.g. 'patch_shape = [1, 512, 512]'.
 53        batch_size: The batch size for training.
 54        n_samples: The number of samples per epoch. By default this will be estimated
 55            based on the patch_shape and size of the volumes used for training.
 56        add_boundary_transform: Whether to add a boundary channel to the training data.
 57        label_dtype: The datatype of the labels returned by the dataloader.
 58        rois: Optional region of interests for training.
 59        sampler: Optional sampler to accept or reject patches for training. 
 60            By default a minimum instance sampler will be used, pass `False` to disable.
 61        ignore_label: Ignore label in the ground-truth. The areas marked by this label will be
 62            ignored in the loss computation. By default this option is not used.
 63        label_transform: Label transform that is applied to the segmentation to compute the targets.
 64            If no label transform is passed (the default) a boundary transform is used.
 65        label_paths: Optional paths containing the labels / annotations for training.
 66            If not given, the labels are expected to be contained in the `data_paths`.
 67        raw_transform: Normalization applied to the raw data. By default the torch_em standardization is used.
 68        loader_kwargs: Additional keyword arguments for the dataloader.
 69
 70    Returns:
 71        The PyTorch dataloader.
 72    """
 73    _, ndim = _determine_ndim(patch_shape)
 74    if label_transform is not None:  # A specific label transform was passed, do nothing.
 75        pass
 76    elif add_boundary_transform:
 77        if ignore_label is None:
 78            label_transform = torch_em.transform.BoundaryTransform(add_binary_target=True)
 79        else:
 80            label_transform = torch_em.transform.label.BoundaryTransformWithIgnoreLabel(
 81                add_binary_target=True, ignore_label=ignore_label
 82            )
 83
 84    else:
 85        if ignore_label is not None:
 86            raise NotImplementedError
 87        label_transform = torch_em.transform.label.connected_components
 88
 89    if ndim == 2:
 90        adjusted_patch_shape = _adjust_patch_shape(ndim, patch_shape)
 91        transform = torch_em.transform.Compose(
 92            torch_em.transform.PadIfNecessary(adjusted_patch_shape), torch_em.transform.get_augmentations(2)
 93        )
 94    else:
 95        transform = torch_em.transform.Compose(
 96            torch_em.transform.PadIfNecessary(patch_shape), torch_em.transform.get_augmentations(3)
 97        )
 98
 99    num_workers = loader_kwargs.pop("num_workers", 4 * batch_size)
100    shuffle = loader_kwargs.pop("shuffle", True)
101
102    if sampler is None:
103        sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=4)
104    elif sampler is False:
105        sampler = None
106
107    if label_paths is None:
108        label_paths = data_paths
109    elif len(label_paths) != len(data_paths):
110        raise ValueError(f"Data paths and label paths don't match: {len(data_paths)} != {len(label_paths)}")
111
112    loader = torch_em.default_segmentation_loader(
113        data_paths, raw_key,
114        label_paths, label_key, sampler=sampler,
115        batch_size=batch_size, patch_shape=patch_shape, ndim=ndim,
116        is_seg_dataset=True, label_transform=label_transform, transform=transform,
117        num_workers=num_workers, shuffle=shuffle, n_samples=n_samples,
118        label_dtype=label_dtype, rois=rois, raw_transform=raw_transform, **loader_kwargs,
119    )
120    return loader
121
122
123def get_unsupervised_loader(
124    data_paths: Tuple[str],
125    raw_key: str,
126    patch_shape: Tuple[int, int, int],
127    batch_size: int,
128    n_samples: Optional[int],
129    sample_mask_paths: Optional[Tuple[str]] = None,
130    sample_mask_key: Optional[str] = None,
131    bg_mask_paths: Optional[Tuple[str]] = None,
132    bg_mask_key: Optional[str] = None,
133    sampler: Optional[callable] = None,
134    exclude_top_and_bottom: bool = False,
135    raw_transform: Optional[callable] = None,
136) -> torch.utils.data.DataLoader:
137    """Get a dataloader for unsupervised segmentation training.
138
139    Args:
140        data_paths: The filepaths to the hdf5 files containing the training data.
141        raw_key: The key that holds the raw data inside of the hdf5.
142        patch_shape: The patch shape used for a training example.
143            In order to run 2d training pass a patch shape with a singleton in the z-axis,
144            e.g. 'patch_shape = [1, 512, 512]'.
145        batch_size: The batch size for training.
146        n_samples: The number of samples per epoch. By default this will be estimated
147            based on the patch_shape and size of the volumes used for training.
148        sample_mask_paths: The filepaths to the corresponding sample masks for each tomogram.
149        sample_mask_key: The key to the sample mask dataset inside each file.
150        bg_mask_paths: The filepaths to the background masks for each tomogram.
151        bg_mask_key: The key to the background mask dataset inside each file.
152        sampler: Optional sampler to accept or reject patches for training. 
153        exclude_top_and_bottom: Whether to exclude the five top and bottom slices to
154            avoid artifacts at the border of tomograms.
155        raw_transform: Normalization applied to the raw data. By default the torch_em standardization is used.
156
157    Returns:
158        The PyTorch dataloader.
159    """
160    if exclude_top_and_bottom:
161        roi = (slice(5, -5), slice(None), slice(None))
162    else:
163        roi = None
164
165    if sample_mask_paths is not None:
166        assert len(data_paths) == len(sample_mask_paths), \
167            f"Expected equal number of data_paths and sample_mask_paths, got {len(data_paths)} and {len(sample_mask_paths)}."
168    if bg_mask_paths is not None:
169        assert len(data_paths) == len(bg_mask_paths), \
170            f"Expected equal number of data_paths and bg_mask_paths, got {len(data_paths)} and {len(bg_mask_paths)}."
171
172    _, ndim = _determine_ndim(patch_shape)
173    if raw_transform is None:
174        raw_transform = torch_em.transform.get_raw_transform()
175    transform = torch_em.transform.get_augmentations(ndim=ndim)
176
177    if n_samples is None:
178        n_samples_per_ds = None
179    else:
180        n_samples_per_ds = int(n_samples / len(data_paths))
181
182    datasets = [
183        RawDatasetWithMasks(
184            raw_path=data_path,
185            raw_key=raw_key,
186            patch_shape=patch_shape,
187            raw_transform=raw_transform,
188            transform=transform,
189            roi=roi,
190            n_samples=n_samples_per_ds,
191            sampler=sampler,
192            ndim=ndim,
193            augmentations=None,
194            sample_mask_path=sample_mask_paths[i] if sample_mask_paths is not None else None,
195            sample_mask_key=sample_mask_key,
196            bg_mask_path=bg_mask_paths[i] if bg_mask_paths is not None else None,
197            bg_mask_key=bg_mask_key,
198        )
199        for i, data_path in enumerate(data_paths)
200    ]
201    ds = torch.utils.data.ConcatDataset(datasets)
202
203    num_workers = 4 * batch_size
204    loader = torch_em.segmentation.get_data_loader(ds, batch_size=batch_size,
205                                                   num_workers=num_workers, shuffle=True)
206    return loader