synapse_net.training.supervised_training

  1import os
  2from glob import glob
  3from typing import Optional, Tuple, Union
  4
  5import torch
  6import torch_em
  7from sklearn.model_selection import train_test_split
  8
  9from synapse_net.inference.inference import get_model_path, get_available_models
 10from synapse_net.inference.util import _restore_legacy_attributes
 11from synapse_net.training.models import get_raw_transform, _check_backbone, _get_model
 12from synapse_net.training.dataloaders import get_supervised_loader
 13
 14
 15def supervised_training(
 16    name: str,
 17    train_paths: Tuple[str],
 18    val_paths: Tuple[str],
 19    label_key: str,
 20    patch_shape: Tuple[int, int, int],
 21    save_root: Optional[str] = None,
 22    raw_key: str = "raw",
 23    batch_size: int = 1,
 24    lr: float = 1e-4,
 25    n_iterations: int = int(1e5),
 26    train_label_paths: Optional[Tuple[str]] = None,
 27    val_label_paths: Optional[Tuple[str]] = None,
 28    train_rois: Optional[Tuple[Tuple[slice]]] = None,
 29    val_rois: Optional[Tuple[Tuple[slice]]] = None,
 30    sampler: Optional[Union[callable, bool]] = None,
 31    n_samples_train: Optional[int] = None,
 32    n_samples_val: Optional[int] = None,
 33    check: bool = False,
 34    ignore_label: Optional[int] = None,
 35    raw_transform: Optional[callable] = None,
 36    label_transform: Optional[callable] = None,
 37    loss_fn: Optional[torch.nn.Module] = None,
 38    in_channels: int = 1,
 39    out_channels: int = 2,
 40    mask_channel: bool = False,
 41    backbone: Optional[str] = None,
 42    model_type: Optional[str] = None,
 43    encoder_checkpoint: Optional[str] = None,
 44    checkpoint_path: Optional[str] = None,
 45    save_every_kth_epoch: Optional[int] = None,
 46    **loader_kwargs,
 47):
 48    """Run supervised segmentation training.
 49
 50    This function trains a UNet for predicting outputs for segmentation.
 51    Expects instance labels and converts them to boundary targets.
 52    This behaviour can be changed by passing custom arguments for `label_transform`
 53    and/or `out_channels`.
 54
 55    Args:
 56        name: The name for the checkpoint to be trained.
 57        train_paths: Filepaths to the hdf5 files for the training data.
 58        val_paths: Filepaths to the df5 files for the validation data.
 59        label_key: The key that holds the labels inside of the hdf5.
 60        patch_shape: The patch shape used for a training example.
 61            In order to run 2d training pass a patch shape with a singleton in the z-axis,
 62            e.g. 'patch_shape = [1, 512, 512]'.
 63        save_root: Folder where the checkpoint will be saved.
 64        raw_key: The key that holds the raw data inside of the hdf5.
 65        batch_size: The batch size for training.
 66        lr: The initial learning rate.
 67        n_iterations: The number of iterations to train for.
 68        train_label_paths: Optional paths containing the label data for training.
 69            If not given, the labels are expected to be part of `train_paths`.
 70        val_label_paths: Optional paths containing the label data for validation.
 71            If not given, the labels are expected to be part of `val_paths`.
 72        train_rois: Optional region of interests for training.
 73        val_rois: Optional region of interests for validation.
 74        sampler: Optional sampler for selecting patches for training.
 75            By default a minimum instance sampler will be used.
 76        n_samples_train: The number of train samples per epoch. By default this will be estimated
 77            based on the patch_shape and size of the volumes used for training.
 78        n_samples_val: The number of val samples per epoch. By default this will be estimated
 79            based on the patch_shape and size of the volumes used for validation.
 80        check: Whether to check the training and validation loaders instead of running training.
 81        ignore_label: Ignore label in the ground-truth. The areas marked by this label will be
 82            ignored in the loss computation. By default this option is not used.
 83        raw_transform: Normalization applied to the raw data. By default the transform for `backbone` is used.
 84            When `backbone` is None, torch_em standardization is used.
 85        label_transform: Label transform that is applied to the segmentation to compute the targets.
 86            If no label transform is passed (the default) a boundary transform is used.
 87        loss_fn: Custom loss function. If None, will default to `torch_em.loss.DiceLoss`.
 88        out_channels: The number of output channels of the UNet.
 89        mask_channel: Whether the last channels in the labels should be used for masking the loss.
 90            This can be used to implement more complex masking operations and is not compatible with `ignore_label`.
 91        backbone: The pretrained encoder. Options: "sam", "sam2", "dinov2", "dinov3" (ViT encoders of a UNETR),
 92            or "cryosiam" (3D CNN encoder of a CryoZooUNet).
 93            The ViT backbones must be set together with `model_type`.
 94        model_type: Model type for the selected ViT `backbone` model family, for example "vit_b" or
 95            "vit_b_em_organelles" for "sam".
 96        encoder_checkpoint: Path to the pretrained encoder weights for `backbone`.
 97            By default they are fetched for `backbone` and `model_type`.
 98        checkpoint_path: Path to the directory with 'best.pt'; continue training this model.
 99        save_every_kth_epoch: Save checkpoints after every kth epoch in a separate file.
100            The corresponding checkpoints will be saved with the naming scheme 'epoch-{epoch}.pt'.
101        loader_kwargs: Additional keyword arguments for the dataloader.
102    """
103    _check_backbone(backbone, model_type, patch_shape)
104
105    if raw_transform is None:
106        raw_transform = get_raw_transform(backbone)[0]
107
108    train_loader = get_supervised_loader(
109        train_paths, raw_key, label_key, patch_shape, batch_size,
110        n_samples=n_samples_train, rois=train_rois, sampler=sampler, ignore_label=ignore_label,
111        label_transform=label_transform, label_paths=train_label_paths, raw_transform=raw_transform,
112        **loader_kwargs,
113    )
114    val_loader = get_supervised_loader(
115        val_paths, raw_key, label_key, patch_shape, batch_size,
116        n_samples=n_samples_val, rois=val_rois, sampler=sampler, ignore_label=ignore_label,
117        label_transform=label_transform, label_paths=val_label_paths, raw_transform=raw_transform,
118        **loader_kwargs,
119    )
120
121    if check:
122        from torch_em.util.debug import check_loader
123        check_loader(train_loader, n_samples=4)
124        check_loader(val_loader, n_samples=4)
125        return
126
127    if checkpoint_path is not None:
128        model = _restore_legacy_attributes(torch_em.util.load_model(checkpoint=checkpoint_path))
129    else:
130        model = _get_model(
131            patch_shape, out_channels, in_channels=in_channels, backbone=backbone,
132            model_type=model_type, encoder_checkpoint=encoder_checkpoint,
133        )
134
135    base_loss = loss_fn if loss_fn is not None else torch_em.loss.DiceLoss()
136    metric = base_loss
137
138    # No ignore label -> we can use default loss.
139    if ignore_label is None and not mask_channel:
140        loss = base_loss
141
142    # If we have an ignore label the loss and metric have to be modified
143    # so that the ignore mask is not used in the gradient calculation.
144    elif ignore_label is not None:
145        loss = torch_em.loss.LossWrapper(
146            loss=base_loss,
147            transform=torch_em.loss.wrapper.MaskIgnoreLabel(
148                ignore_label=ignore_label, masking_method="multiply",
149            )
150        )
151        metric = loss
152    elif mask_channel:
153        loss = torch_em.loss.LossWrapper(
154            loss=base_loss,
155            transform=torch_em.loss.wrapper.ApplyAndRemoveMask(
156                masking_method="crop" if out_channels == 1 else "multiply")
157        )
158        metric = loss
159    else:
160        loss = base_loss
161        metric = loss
162
163    trainer = torch_em.default_segmentation_trainer(
164        name=name,
165        model=model,
166        train_loader=train_loader,
167        val_loader=val_loader,
168        learning_rate=lr,
169        mixed_precision=True,
170        log_image_interval=100,
171        compile_model=False,
172        save_root=save_root,
173        loss=loss,
174        metric=metric,
175    )
176    trainer.fit(n_iterations, save_every_kth_epoch=save_every_kth_epoch)
177
178
179def _derive_key_from_files(files, key):
180    # Get all file extensions (general wild-cards may pick up files with multiple extensions).
181    extensions = list(set([os.path.splitext(ff)[1] for ff in files]))
182
183    # If we have more than 1 file extension we just use the key that was passed,
184    # as it is unclear how to derive a consistent key.
185    if len(extensions) > 1:
186        return files, key
187
188    ext = extensions[0]
189    extension_to_key = {".tif": None, ".mrc": "data", ".rec": "data"}
190
191    # Derive the key from the extension if the key is None.
192    if key is None and ext in extension_to_key:
193        key = extension_to_key[ext]
194    # If the key is None and can't be derived raise an error.
195    elif key is None and ext not in extension_to_key:
196        raise ValueError(
197            f"You have not passed a key for the data in {ext} format, for which the key cannot be derived."
198        )
199    # If the key was passed and doesn't match the extension raise an error.
200    elif key is not None and ext in extension_to_key and key != extension_to_key[ext]:
201        raise ValueError(
202            f"The expected key {extension_to_key[ext]} for format {ext} did not match the passed key {key}."
203        )
204    return files, key
205
206
207def _parse_input_folder(folder, pattern, key):
208    files = sorted(glob(os.path.join(folder, "**", pattern), recursive=True))
209    return _derive_key_from_files(files, key)
210
211
212def _parse_input_files(args):
213    train_image_paths, raw_key = _parse_input_folder(args.train_folder, args.image_file_pattern, args.raw_key)
214    train_label_paths, label_key = _parse_input_folder(args.label_folder, args.label_file_pattern, args.label_key)
215    if len(train_image_paths) != len(train_label_paths):
216        raise ValueError(
217            f"The image and label paths parsed from {args.train_folder} and {args.label_folder} don't match."
218            f"The image folder contains {len(train_image_paths)}, the label folder contains {len(train_label_paths)}."
219        )
220
221    if args.val_folder is None:
222        if args.val_label_folder is not None:
223            raise ValueError("You have passed a val_label_folder, but not a val_folder.")
224        train_image_paths, val_image_paths, train_label_paths, val_label_paths = train_test_split(
225            train_image_paths, train_label_paths, test_size=args.val_fraction, random_state=42
226        )
227    else:
228        if args.val_label_folder is None:
229            raise ValueError("You have passed a val_folder, but not a val_label_folder.")
230        val_image_paths, _ = _parse_input_folder(args.val_folder, args.image_file_pattern, raw_key)
231        val_label_paths, _ = _parse_input_folder(args.val_label_folder, args.label_file_pattern, label_key)
232
233    return train_image_paths, train_label_paths, val_image_paths, val_label_paths, raw_key, label_key
234
235
236def _parse_checkpoint(initial_model):
237    if initial_model is None:
238        return None
239    if os.path.exists(initial_model):
240        return initial_model
241    model_path = get_model_path(initial_model)
242    return model_path
243
244
245def main():
246    """@private
247    """
248    import argparse
249
250    parser = argparse.ArgumentParser(
251        description="Train a model for foreground and boundary segmentation via supervised learning.\n\n"
252        "You can use this function to train a model for vesicle segmentation, or another segmentation task, like this:\n"  # noqa
253        "synapse_net.run_supervised_training -n my_model -i /path/to/images -l /path/to/labels --patch_shape 32 192 192\n"  # noqa
254        "The trained model will be saved in the folder 'checkpoints/my_model' (or whichever name you pass to the '-n' argument)."  # noqa
255        "You can then use this model for segmentation with the SynapseNet GUI or CLI. "
256        "Check out the information below for details on the arguments of this function.",
257        formatter_class=argparse.RawTextHelpFormatter
258    )
259    parser.add_argument("-n", "--name", required=True, help="The name of the model to be trained.")
260    parser.add_argument("-p", "--patch_shape", nargs=3, type=int, help="The patch shape for training.")
261
262    # Folders with training data, containing raw/image data and labels.
263    parser.add_argument("-i", "--train_folder", required=True, help="The input folder with the training image data.")
264    parser.add_argument("--image_file_pattern", default="*",
265                        help="The pattern for selecting image files. For example, '*.mrc' to select all mrc files.")
266    parser.add_argument("--raw_key",
267                        help="The internal path for the raw data. If not given, will be determined based on the file extension.")  # noqa
268    parser.add_argument("-l", "--label_folder", required=True, help="The input folder with the training labels.")
269    parser.add_argument("--label_file_pattern", default="*",
270                        help="The pattern for selecting label files. For example, '*.tif' to select all tif files.")
271    parser.add_argument("--label_key",
272                        help="The internal path for the label data. If not given, will be determined based on the file extension.")  # noqa
273
274    # Optional folders with validation data. If not given the training data is split into train/val.
275    parser.add_argument("--val_folder",
276                        help="The input folder with the validation data. If not given the training data will be split for validation")  # noqa
277    parser.add_argument("--val_label_folder",
278                        help="The input folder with the validation labels. If not given the training data will be split for validation.")  # noqa
279
280    # Optional: choose a model for initializing the weights.
281    available_models = get_available_models()
282    parser.add_argument(
283        "--initial_model",
284        help="Choose a model checkpoint for weight initialization.\n"
285        "This may either be the path to an existing model checkpoint or the name of a pretrained model.\n"
286        f"The following pretrained models are available: {available_models}.\n"
287        "If not given, the model will be randomly initialized."
288    )
289
290    # More optional argument:
291    parser.add_argument("--batch_size", type=int, default=1, help="The batch size for training.")
292    parser.add_argument("--n_samples_train", type=int, help="The number of samples per epoch for training. If not given will be derived from the data size.")  # noqa
293    parser.add_argument("--n_samples_val", type=int, help="The number of samples per epoch for validation. If not given will be derived from the data size.")  # noqa
294    parser.add_argument("--val_fraction", type=float, default=0.15, help="The fraction of the data to use for validation. This has no effect if 'val_folder' and 'val_label_folder' were passed.")  # noqa
295    parser.add_argument("--check", action="store_true", help="Visualize samples from the data loaders to ensure correct data instead of running training.")  # noqa
296    parser.add_argument("--n_iterations", type=int, default=int(1e5), help="The maximal number of iterations to train for.")  # noqa
297    parser.add_argument("--save_root", help="Root path for saving the checkpoint and log dir.")
298    args = parser.parse_args()
299
300    train_image_paths, train_label_paths, val_image_paths, val_label_paths, raw_key, label_key =\
301        _parse_input_files(args)
302    checkpoint_path = _parse_checkpoint(args.initial_model)
303
304    supervised_training(
305        name=args.name, train_paths=train_image_paths, val_paths=val_image_paths,
306        train_label_paths=train_label_paths, val_label_paths=val_label_paths,
307        raw_key=raw_key, label_key=label_key, patch_shape=args.patch_shape, batch_size=args.batch_size,
308        n_samples_train=args.n_samples_train, n_samples_val=args.n_samples_val,
309        check=args.check, n_iterations=args.n_iterations, save_root=args.save_root,
310        checkpoint_path=checkpoint_path
311    )