synapse_net.training.models

  1"""Model factories and pretrained backbone utilities for training."""
  2import warnings
  3from typing import Tuple
  4import numpy as np
  5
  6import torch
  7
  8from torch_em.model import UNet2d, AnisotropicUNet
  9from torch_em.model.cryozoo import get_cryozoounet
 10from torch_em.model.unetr import get_unetr_model
 11from torch_em.transform.raw import normalize_percentile, get_raw_transform_for_pretrained_vit
 12
 13from .dataloaders import _determine_ndim
 14
 15try:
 16    from micro_sam.v1.util import models as microsam_models
 17except ImportError:
 18    microsam_models = None
 19
 20try:
 21    from micro_sam.v2.util import models as microsam2_models
 22except ImportError:
 23    microsam2_models = None
 24
 25try:
 26    from micro_sam.models.vfm import get_vfm_model
 27except ImportError:
 28    get_vfm_model = None
 29
 30
 31def get_3d_model(
 32    out_channels: int,
 33    in_channels: int = 1,
 34    scale_factors: Tuple[Tuple[int, int, int]] = [[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
 35    initial_features: int = 32,
 36    final_activation: str = "Sigmoid",
 37) -> torch.nn.Module:
 38    """Get the U-Net model for 3D segmentation tasks.
 39
 40    Args:
 41        out_channels: The number of output channels of the network.
 42        scale_factors: The downscaling factors for each level of the U-Net encoder.
 43        initial_features: The number of features in the first level of the U-Net.
 44            The number of features increases by a factor of two in each level.
 45        final_activation: The activation applied to the last output layer.
 46
 47    Returns:
 48        The U-Net.
 49    """
 50    model = AnisotropicUNet(
 51        scale_factors=scale_factors,
 52        in_channels=in_channels,
 53        out_channels=out_channels,
 54        initial_features=initial_features,
 55        gain=2,
 56        final_activation=final_activation,
 57    )
 58    return model
 59
 60
 61def get_2d_model(
 62    out_channels: int,
 63    in_channels: int = 1,
 64    initial_features: int = 32,
 65    final_activation: str = "Sigmoid",
 66) -> torch.nn.Module:
 67    """Get the U-Net model for 2D segmentation tasks.
 68
 69    Args:
 70        out_channels: The number of output channels of the network.
 71        initial_features: The number of features in the first level of the U-Net.
 72            The number of features increases by a factor of two in each level.
 73        final_activation: The activation applied to the last output layer.
 74
 75    Returns:
 76        The U-Net.
 77    """
 78    model = UNet2d(
 79        in_channels=in_channels,
 80        out_channels=out_channels,
 81        initial_features=initial_features,
 82        gain=2,
 83        depth=4,
 84        final_activation=final_activation,
 85    )
 86    return model
 87
 88
 89def _check_backbone(backbone, model_type, patch_shape):
 90    if backbone == "cryosiam":
 91        if model_type is not None:
 92            raise ValueError(f"The 'cryosiam' backbone has no model_type; got model_type='{model_type}'.")
 93        if _determine_ndim(patch_shape)[0]:
 94            raise ValueError(f"The 'cryosiam' backbone only supports 3D training; got patch_shape={patch_shape}.")
 95        if any(s % 8 for s in patch_shape):
 96            raise ValueError(
 97                f"The patch_shape for 'cryosiam' backbone must be divisible by 8; got patch_shape={patch_shape}."
 98            )
 99    elif (backbone is None) != (model_type is None):
100        raise ValueError(
101            "backbone and model_type must both be set for ViT backbones; "
102            f"got backbone={backbone} and model_type={model_type}."
103        )
104
105
106def _get_model(
107    patch_shape, out_channels, in_channels=1, backbone=None, model_type=None,
108    encoder_checkpoint=None, final_activation="Sigmoid",
109):
110    is_2d, ndim = _determine_ndim(patch_shape)
111
112    if backbone == "cryosiam":
113        assert in_channels == 1
114        return get_cryozoounet(out_channels=out_channels, encoder=backbone, final_activation=final_activation)
115
116    if backbone is not None:
117        assert in_channels == 1
118        if encoder_checkpoint is None:
119            encoder_checkpoint = _get_checkpoint(backbone=backbone, model_type=model_type)
120        return get_unetr_model(
121            ndim=ndim, backbone=backbone, model_type=model_type,
122            encoder_checkpoint=encoder_checkpoint,
123            out_channels=out_channels, final_activation=final_activation,
124        )
125
126    if is_2d:
127        return get_2d_model(out_channels=out_channels, in_channels=in_channels, final_activation=final_activation)
128
129    return get_3d_model(out_channels=out_channels, in_channels=in_channels, final_activation=final_activation)
130
131
132def normalize_percentile_to_0_1_inverted(raw):
133    """Normalize the raw data with percentiles to [0, 1] and invert the contrast."""
134    raw = 1 - normalize_percentile(raw)
135    return torch.clamp(raw, 0, 1) if torch.is_tensor(raw) else np.clip(raw, 0, 1)
136
137
138def get_raw_transform(backbone):
139    """Get the raw transform and the maximum input value for a pretrained backbone.
140
141    The ViT backbones are handled by `torch_em.transform.raw.get_raw_transform_for_pretrained_vit`.
142
143    Args:
144        backbone: The pretrained encoder. Options: "sam", "sam2", "dinov2", "dinov3", "cryosiam", or None.
145
146    Returns:
147        The raw transform. It normalizes the raw data with percentiles to [0, 255] for "sam",
148            and to [0, 1] for "sam2", "dinov2" and "dinov3". For "cryosiam" it also inverts the contrast.
149            None if `backbone` is None, so the loaders use the default torch-em standardization.
150        The maximum input value `clip_max`. The intensity augmentations clip their output to
151            [0, clip_max], so that the input stays in the expected range. None if `backbone` is None.
152    """
153    if backbone == "cryosiam":
154        return normalize_percentile_to_0_1_inverted, 1
155
156    if backbone not in ("sam", "sam2", "dinov2", "dinov3", None):
157        raise ValueError(f"Unsupported backbone '{backbone}'.")
158    return get_raw_transform_for_pretrained_vit(backbone)
159
160
161def _get_checkpoint(backbone, model_type, return_decoder_path=False):
162
163    if backbone == "sam":
164        if microsam_models is None:
165            raise RuntimeError(
166                "The 'sam' backbone requires micro_sam. Install it with 'conda install -c conda-forge micro_sam'."
167                )
168        model_registry = microsam_models()
169        encoder_checkpoint = model_registry.fetch(model_type, progressbar=True)
170
171        if return_decoder_path:
172            decoder_name = f"{model_type}_decoder"
173            decoder_path = model_registry.fetch(
174                decoder_name, progressbar=True
175            ) if decoder_name in model_registry.registry else None
176
177    elif backbone == "sam2":
178        if microsam2_models is None:
179            raise RuntimeError(
180                "The 'sam2' backbone requires micro_sam with SAM2 support ('micro_sam.v2') and sam2."
181                )
182        model_registry = microsam2_models()
183        encoder_checkpoint = model_registry.fetch(model_type, progressbar=True)
184
185        if return_decoder_path:
186            decoder_name = f"{model_type}_decoder"
187            decoder_path = model_registry.fetch(
188                decoder_name, progressbar=True
189            ) if decoder_name in model_registry.registry else None
190
191    elif backbone == "dinov2":
192        if get_vfm_model is None:
193            raise RuntimeError("The 'dinov2' backbone requires micro_sam with the 'micro_sam.models.vfm' module.")
194        encoder_checkpoint = get_vfm_model(f"{model_type}_dinov2", device="cpu").model.state_dict()
195        decoder_path = None
196
197    elif backbone == "dinov3":
198        if get_vfm_model is None:
199            raise RuntimeError("The 'dinov3' backbone requires micro_sam with the 'micro_sam.models.vfm' module.")
200        encoder_checkpoint = get_vfm_model(f"{model_type}_dinov3", device="cpu").model.state_dict()
201        decoder_path = None
202    else:
203        raise ValueError(f"Unsupported backbone '{backbone}'.")
204
205    if return_decoder_path:
206        return encoder_checkpoint, decoder_path
207    else:
208        return encoder_checkpoint
209
210
211def _init_microsam_decoder(model, backbone, model_type, out_channels):
212    # Get the micro-sam decoder path
213    _, decoder_path = _get_checkpoint(backbone=backbone, model_type=model_type, return_decoder_path=True)
214
215    if decoder_path is None:
216        warnings.warn("Oopsie daisy. You choose a `micro-sam` or model without pretrained decoder weights.")
217        return model
218
219    # Load the pretrained decoder weights.
220    decoder_state = torch.load(decoder_path, map_location="cpu", weights_only=True)
221
222    # Super hacky way of initializing decoder weights - please don't try this at home.
223    unetr_state_dict = model.state_dict()
224    for k, v in unetr_state_dict.items():
225        # If the key starts with encoder, we don't care about it.
226        if k.startswith("encoder"):
227            continue
228
229        # If the 'out_channels' are different than expected, i.e. 3, we say ciao ciao to it.
230        if out_channels != 3 and k.startswith("out_conv"):
231            warnings.warn(f"Seems like output channels are different than 'micro-sam'. Hence, we reinitialize '{k}'.")
232            unetr_state_dict[k] = v
233            continue
234
235        # If it's some weird parameters, we don't care about them too.
236        if k not in decoder_state:
237            warnings.warn(f"Could not find '{k}' in the pretrained decoder state dict. Hence, we reinitialize '{k}.")
238            unetr_state_dict[k] = v
239            continue
240
241        # If the parameter shapes mismatch, we don't care about them either.
242        if decoder_state[k].shape != v.shape:
243            warnings.warn(f"There's a shape mismatch for '{k}'. Hence, we reinitialize '{k}'.")
244            unetr_state_dict[k] = v
245            continue
246
247        # Keep the expected decoder weights.
248        unetr_state_dict[k] = decoder_state[k]
249
250    model.load_state_dict(unetr_state_dict)
251
252    return model