torch_em.model.cryozoo

  1"""CryoZoo: pretrained models for cryo-electron tomography (cryo-ET) in the torch-em segmentation framework.
  2
  3CryoZoo integrates published pretrained cryo-ET models into torch-em, so that you can finetune and use them
  4for segmentation like the other torch-em models.
  5CryoSiam is one of these integrations, and the CryoZooUNet uses its pretrained encoder.
  6
  7CryoSiam is described in https://doi.org/10.1101/2025.11.11.687379.
  8Please cite it if you use this piece of code or pretrained weights in your research.
  9
 10The authors release the weights under the MIT license at https://huggingface.co/frosinastojanovska/cryosiam_v1.0.
 11This module does not contain code from CryoSiam. It builds the encoder from the ResNet of MONAI, like CryoSiam does.
 12Thus, the weights of the CryoSiam checkpoint load without changes.
 13"""
 14
 15import os
 16from collections import OrderedDict
 17from typing import List, Optional, Union
 18
 19import numpy as np
 20
 21import torch
 22import torch.nn as nn
 23
 24try:
 25    from monai.networks.nets import ResNet
 26except ImportError:
 27    ResNet = None
 28
 29from ..util.util import get_cache_directory
 30from .unet import UNetBase, Decoder, ConvBlock3d, Upsampler3d
 31
 32
 33CRYOSIAM_URL = "https://huggingface.co/frosinastojanovska/cryosiam_v1.0/resolve/main/dense_simsiam_pretrained.ckpt"
 34CRYOSIAM_CHECKSUM = "19957ef7eba45658432a81d37ac1764db65f97ddbfcd55bee3c190f51a9b17e7"
 35
 36ENCODERS = {}  # The encoders of the CryoZooUNet by name. The register_encoder decorator fills it at import.
 37
 38
 39def register_encoder(name: str):
 40    """Add an encoder class to `ENCODERS` under the given name.
 41
 42    Use this function as a decorator for the encoder class.
 43
 44    Args:
 45        name: The name of the encoder.
 46
 47    Returns:
 48        The decorator, which adds the encoder class and returns it unchanged.
 49    """
 50    def register(encoder_class):
 51        ENCODERS[name] = encoder_class
 52        return encoder_class
 53    return register
 54
 55
 56class CryoZooEncoder(nn.Module):
 57    """The base class for the CNN encoders of the CryoZooUNet.
 58
 59    To add an encoder, make a subclass and decorate it with `register_encoder`. The subclass must:
 60    - set `features`: the number of channels of each skip connection and of the bottleneck, from fine to coarse.
 61    - set `scale_factors`: the downsampling factor between two consecutive entries of `features`.
 62    - set `stem_scale_factor`: the downsampling factor between the input and the first skip connection.
 63    - implement `forward`: it returns the bottleneck and the list of skip connections, from fine to coarse.
 64
 65    If pretrained weights exist, the subclass must also:
 66    - set `url`, `checksum` and `cache_name`: the URL and the SHA256 checksum of the checkpoint,
 67      and the name of its folder in the torch-em cache directory.
 68    - implement `load_pretrained`.
 69
 70    A scale factor is an integer or a list with one integer for each spatial axis (Z, Y, X).
 71
 72    Args:
 73        in_channels: The number of input channels.
 74    """
 75    url: Optional[str] = None
 76    checksum: Optional[str] = None
 77    cache_name: Optional[str] = None
 78
 79    def __init__(self, in_channels: int = 1):
 80        super().__init__()
 81        self.in_channels = in_channels
 82        self.return_outputs = True
 83        self.features: List[int] = []
 84        self.scale_factors: List[Union[int, List[int]]] = []
 85        self.stem_scale_factor: Union[int, List[int]] = 1
 86
 87    @property
 88    def out_channels(self):
 89        """The number of channels of the bottleneck."""
 90        return self.features[-1]
 91
 92    def __len__(self):
 93        return len(self.scale_factors)
 94
 95    @classmethod
 96    def get_checkpoint(cls, checkpoint_path: Optional[Union[str, os.PathLike]] = None, download: bool = True) -> str:
 97        """Get the filepath to the pretrained checkpoint.
 98
 99        If you pass `checkpoint_path`, this method uses this file.
100        Otherwise, it uses the checkpoint in the torch-em cache directory, and it downloads the checkpoint if needed.
101        See `torch_em.util.util.get_cache_directory` for the cache directory.
102
103        Args:
104            checkpoint_path: The filepath to a checkpoint on disk.
105            download: Whether to download the checkpoint if it is not in the cache yet.
106
107        Returns:
108            The filepath to the checkpoint.
109        """
110        if checkpoint_path is not None:
111            if not os.path.exists(checkpoint_path):
112                raise FileNotFoundError(f"The checkpoint {checkpoint_path} does not exist.")
113            return str(checkpoint_path)
114
115        if cls.url is None:
116            raise NotImplementedError(f"The encoder {cls.__name__} does not have pretrained weights.")
117
118        import pooch  # We import pooch here, so that 'import torch_em' does not need it.
119
120        folder = os.path.join(get_cache_directory(), "models", cls.cache_name)
121        fname = cls.url.split("/")[-1]
122        if not download and not os.path.exists(os.path.join(folder, fname)):
123            raise RuntimeError(f"The checkpoint is not in {folder}. Set download=True to download it.")
124
125        return pooch.retrieve(url=cls.url, known_hash=f"sha256:{cls.checksum}", fname=fname, path=folder)
126
127    def load_pretrained(self, checkpoint: Union[str, os.PathLike, OrderedDict]) -> None:
128        """Load the pretrained weights into the encoder.
129
130        Args:
131            checkpoint: The filepath to a checkpoint or a state dict for the encoder.
132        """
133        raise NotImplementedError(f"The encoder {type(self).__name__} does not have pretrained weights.")
134
135
136def load_cryosiam_encoder_state(checkpoint: Union[str, os.PathLike]) -> OrderedDict:
137    """Load the encoder weights from a CryoSiam checkpoint.
138
139    Args:
140        checkpoint: The filepath to the CryoSiam checkpoint.
141
142    Returns:
143        The state dict for the MONAI ResNet of `CryoSiamEncoder`.
144    """
145    state = torch.load(checkpoint, map_location="cpu", weights_only=True)["state_dict"]
146    prefix = "_model.encoder."
147    # The 'fc' layer belongs to the global projection head of the pretraining.
148    return OrderedDict(
149        (k[len(prefix):], v) for k, v in state.items() if k.startswith(prefix) and not k.startswith(prefix + "fc.")
150    )
151
152
153@register_encoder("cryosiam")
154class CryoSiamEncoder(CryoZooEncoder):
155    """The 3D ResNet encoder of CryoSiam, which uses the ResNet from MONAI.
156
157    It has a 7x7x7 stem without stride and four stages with one basic block each.
158    The first three stages give the skip connections and the last stage gives the bottleneck.
159
160    Args:
161        in_channels: The number of input channels. The pretrained weights expect one channel.
162    """
163    url = CRYOSIAM_URL
164    checksum = CRYOSIAM_CHECKSUM
165    cache_name = "cryosiam"
166
167    def __init__(self, in_channels: int = 1):
168        super().__init__(in_channels=in_channels)
169        if ResNet is None:
170            raise RuntimeError("The CryoSiam encoder needs MONAI. Install it from conda-forge or with pip.")
171
172        self.features = [64, 128, 256, 512]
173        self.scale_factors = [2, 2, 2]
174        self.stem_scale_factor = 1
175
176        # We pin all the settings that the CryoSiam checkpoint depends on, in case the MONAI defaults change.
177        self.resnet = ResNet(
178            block="basic", layers=[1, 1, 1, 1], block_inplanes=self.features, spatial_dims=3,
179            n_input_channels=in_channels, conv1_t_size=7, conv1_t_stride=1, no_max_pool=True, shortcut_type="B",
180            feed_forward=False, bias_downsample=True, act=("relu", {"inplace": False}), norm="batch",
181        )
182
183    def load_pretrained(self, checkpoint: Union[str, os.PathLike, OrderedDict]) -> None:
184        state = checkpoint if isinstance(checkpoint, OrderedDict) else load_cryosiam_encoder_state(checkpoint)
185        self.resnet.load_state_dict(state)
186
187    def forward(self, x):
188        """Apply the encoder to the input data.
189
190        Args:
191            x: The input tensor with the shape (B, C, Z, Y, X).
192
193        Returns:
194            The bottleneck, and the list of skip connections from fine to coarse if `return_outputs` is True.
195        """
196        net = self.resnet
197        x = net.act(net.bn1(net.conv1(x)))
198        encoder_out = []
199        for layer in (net.layer1, net.layer2, net.layer3):
200            x = layer(x)
201            encoder_out.append(x)
202        x = net.layer4(x)
203
204        if self.return_outputs:
205            return x, encoder_out
206        else:
207            return x
208
209
210class OutputHead(nn.Module):
211    """@private
212    """
213    def __init__(self, in_channels, out_channels, scale_factor):
214        super().__init__()
215        self.out_channels = out_channels
216        self.upsampler = Upsampler3d(scale_factor, in_channels, in_channels)
217        self.conv = nn.Conv3d(in_channels, out_channels, 1)
218
219    def forward(self, x):
220        return self.conv(self.upsampler(x))
221
222
223class CryoZooUNet(UNetBase):
224    """A 3D U-Net with a pretrained CNN encoder and a convolutional decoder from torch-em.
225
226    The decoder always uses the skip connections of the encoder.
227    Each spatial input dimension must be divisible by the total downsampling factor of the encoder.
228    For the CryoSiam encoder, this factor is 8.
229    The pretrained CryoSiam weights expect tomograms with inverted contrast and intensities scaled to [0, 1].
230
231    `init_kwargs` does not contain `encoder_checkpoint`. When torch-em restores a trained model,
232    it loads all the weights from the saved training state. Thus it does not load the pretrained weights again.
233
234    Args:
235        out_channels: The number of output channels.
236        encoder: The name of the encoder. See `ENCODERS` for the supported names.
237        in_channels: The number of input channels. The pretrained CryoSiam weights expect one channel.
238        encoder_checkpoint: The filepath to a checkpoint or a state dict for the encoder.
239            The model loads the encoder weights from it. Use `get_cryozoounet` to get the pretrained weights.
240            If you do not pass it, the encoder starts with random weights.
241        final_activation: The activation applied after the output convolution.
242        postprocessing: A postprocessing function to apply after the U-Net output.
243        check_shape: Whether to check the input shape to the U-Net forward call.
244        conv_block_kwargs: The keyword arguments for the convolutional blocks of the decoder.
245    """
246    def __init__(
247        self,
248        out_channels: int,
249        encoder: str = "cryosiam",
250        in_channels: int = 1,
251        encoder_checkpoint: Optional[Union[str, os.PathLike, OrderedDict]] = None,
252        final_activation: Optional[Union[nn.Module, str]] = None,
253        postprocessing: Optional[Union[nn.Module, str]] = None,
254        check_shape: bool = True,
255        **conv_block_kwargs,
256    ):
257        if encoder not in ENCODERS:
258            raise ValueError(f"The encoder '{encoder}' is not supported. Choose one of {list(ENCODERS)}.")
259        encoder_model = ENCODERS[encoder](in_channels=in_channels)
260
261        features_decoder = encoder_model.features[::-1]
262        decoder = Decoder(
263            features=features_decoder,
264            scale_factors=encoder_model.scale_factors[::-1],
265            conv_block_impl=ConvBlock3d,
266            sampler_impl=Upsampler3d,
267            skip_channels=features_decoder[1:],
268            **conv_block_kwargs,
269        )
270
271        if np.all(np.array(encoder_model.stem_scale_factor) == 1):
272            out_conv = nn.Conv3d(features_decoder[-1], out_channels, 1)
273        else:
274            out_conv = OutputHead(features_decoder[-1], out_channels, encoder_model.stem_scale_factor)
275
276        super().__init__(
277            encoder=encoder_model,
278            base=nn.Identity(),
279            decoder=decoder,
280            out_conv=out_conv,
281            final_activation=final_activation,
282            postprocessing=postprocessing,
283            check_shape=check_shape,
284        )
285        self.init_kwargs = {
286            "out_channels": out_channels, "encoder": encoder, "in_channels": in_channels,
287            "final_activation": final_activation, "postprocessing": postprocessing, "check_shape": check_shape,
288            **conv_block_kwargs,
289        }
290
291        if encoder_checkpoint is not None:
292            self.encoder.load_pretrained(encoder_checkpoint)
293
294    def _check_shape(self, x):
295        spatial_shape = tuple(x.shape)[2:]
296        factor = np.ones(len(spatial_shape), dtype=int) * np.array(self.encoder.stem_scale_factor)
297        for scale_factor in self.encoder.scale_factors:
298            factor = factor * np.array(scale_factor)
299        if any(sh % fac != 0 for sh, fac in zip(spatial_shape, factor)):
300            raise ValueError(f"The input shape {spatial_shape} must be divisible by {factor.tolist()}.")
301
302
303def get_cryozoounet(
304    out_channels: int,
305    encoder: str = "cryosiam",
306    pretrained: bool = True,
307    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
308    **kwargs,
309) -> CryoZooUNet:
310    """Get a CryoZooUNet with a pretrained CNN encoder.
311
312    Args:
313        out_channels: The number of output channels.
314        encoder: The name of the encoder. See `ENCODERS` for the supported names.
315        pretrained: Whether to initialize the encoder with its pretrained weights.
316        checkpoint_path: The filepath to a checkpoint on disk.
317            If you do not pass it, the function downloads the pretrained weights to the torch-em cache directory
318            when you use them the first time.
319        kwargs: Additional keyword arguments for `CryoZooUNet`.
320
321    Returns:
322        The CryoZooUNet.
323    """
324    if encoder not in ENCODERS:
325        raise ValueError(f"The encoder '{encoder}' is not supported. Choose one of {list(ENCODERS)}.")
326    encoder_checkpoint = ENCODERS[encoder].get_checkpoint(checkpoint_path) if pretrained else None
327    return CryoZooUNet(out_channels=out_channels, encoder=encoder, encoder_checkpoint=encoder_checkpoint, **kwargs)