Skip to content

CellVitPanoptic

Bases: BaseModelPanoptic

Source code in src/histolytics/models/cellvit_panoptic.py 
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
class CellVitPanoptic(BaseModelPanoptic):
    model_name = "cellvit_panoptic"

    def __init__(
        self,
        n_nuc_classes: int,
        n_tissue_classes: int,
        enc_name: str = "samvit_base_patch16",
        enc_pretrain: bool = True,
        enc_freeze: bool = False,
        device: torch.device = torch.device("cuda"),
        model_kwargs: Dict[str, Any] = {},
    ) -> None:
        """CellVitPanoptic model for panoptic segmentation of nuclei and tissues.

        Note:
            [CellVit article](https://arxiv.org/abs/2306.15350)

        Parameters:
            n_nuc_classes (int):
                Number of nuclei type classes.
            n_tissue_classes (int):
                Number of tissue type classes.
            enc_name (str):
                Name of the pytorch-image-models encoder.
            enc_pretrain (bool):
                Whether to use pretrained weights in the encoder.
            enc_freeze (bool):
                Freeze encoder weights for training.
            device (torch.device):
                Device to run the model on.
            model_kwargs (dict):
                Additional keyword arguments for the model.
        """
        super().__init__()
        self.model = cellvit_panoptic(
            n_nuc_classes=n_nuc_classes,
            n_tissue_classes=n_tissue_classes,
            enc_name=enc_name,
            enc_pretrain=enc_pretrain,
            enc_freeze=enc_freeze,
            **model_kwargs,
        )

        self.device = device
        self.model.to(device)

    def set_inference_mode(self, mixed_precision: bool = True) -> None:
        """Set model to inference mode."""
        self.model.eval()
        self.predictor = Predictor(
            model=self.model,
            mixed_precision=mixed_precision,
        )
        self.post_processor = PostProcessor(postproc_method="hovernet")
        self.inference_mode = True

__init__ 

__init__(n_nuc_classes: int, n_tissue_classes: int, enc_name: str = 'samvit_base_patch16', enc_pretrain: bool = True, enc_freeze: bool = False, device: device = torch.device('cuda'), model_kwargs: Dict[str, Any] = {}) -> None

CellVitPanoptic model for panoptic segmentation of nuclei and tissues.

Note

CellVit article

Parameters:

Name Type Description Default
n_nuc_classes int

Number of nuclei type classes.

 required
n_tissue_classes int

Number of tissue type classes.

 required
enc_name str

Name of the pytorch-image-models encoder.

 'samvit_base_patch16'
enc_pretrain bool

Whether to use pretrained weights in the encoder.

 True
enc_freeze bool

Freeze encoder weights for training.

 False
device device

Device to run the model on.

 device('cuda')
model_kwargs dict

Additional keyword arguments for the model.

 {}
Source code in src/histolytics/models/cellvit_panoptic.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def __init__(
    self,
    n_nuc_classes: int,
    n_tissue_classes: int,
    enc_name: str = "samvit_base_patch16",
    enc_pretrain: bool = True,
    enc_freeze: bool = False,
    device: torch.device = torch.device("cuda"),
    model_kwargs: Dict[str, Any] = {},
) -> None:
    """CellVitPanoptic model for panoptic segmentation of nuclei and tissues.

    Note:
        [CellVit article](https://arxiv.org/abs/2306.15350)

    Parameters:
        n_nuc_classes (int):
            Number of nuclei type classes.
        n_tissue_classes (int):
            Number of tissue type classes.
        enc_name (str):
            Name of the pytorch-image-models encoder.
        enc_pretrain (bool):
            Whether to use pretrained weights in the encoder.
        enc_freeze (bool):
            Freeze encoder weights for training.
        device (torch.device):
            Device to run the model on.
        model_kwargs (dict):
            Additional keyword arguments for the model.
    """
    super().__init__()
    self.model = cellvit_panoptic(
        n_nuc_classes=n_nuc_classes,
        n_tissue_classes=n_tissue_classes,
        enc_name=enc_name,
        enc_pretrain=enc_pretrain,
        enc_freeze=enc_freeze,
        **model_kwargs,
    )

    self.device = device
    self.model.to(device)

set_inference_mode 

set_inference_mode(mixed_precision: bool = True) -> None

Set model to inference mode.

Source code in src/histolytics/models/cellvit_panoptic.py 
 97
 98
 99
100
101
102
103
104
105
def set_inference_mode(self, mixed_precision: bool = True) -> None:
    """Set model to inference mode."""
    self.model.eval()
    self.predictor = Predictor(
        model=self.model,
        mixed_precision=mixed_precision,
    )
    self.post_processor = PostProcessor(postproc_method="hovernet")
    self.inference_mode = True

from_pretrained classmethod 

from_pretrained(weights: Union[str, Path], device: device = torch.device('cuda'), model_kwargs: Dict[str, Any] = {})

Load the model from pretrained weights.

Parameters:

Name Type Description Default
model_name str

Name of the pretrained model.

 required
device device

Device to run the model on.

 device('cuda')
model_kwargs Dict[str, Any]

Additional arguments for the model.

 {}

Examples:

>>> model = Model.from_pretrained(<str or Path to weights>, device=torch.device("cuda"))
Source code in src/histolytics/models/_base_model.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@classmethod
def from_pretrained(
    cls,
    weights: Union[str, Path],
    device: torch.device = torch.device("cuda"),
    model_kwargs: Dict[str, Any] = {},
):
    """Load the model from pretrained weights.

    Parameters:
        model_name (str):
            Name of the pretrained model.
        device (torch.device):
            Device to run the model on.
        model_kwargs (Dict[str, Any]):
            Additional arguments for the model.

    Examples:
        >>> model = Model.from_pretrained(<str or Path to weights>, device=torch.device("cuda"))
    """
    weights_path = Path(weights)
    if not weights_path.is_file():
        if weights_path.as_posix() in PRETRAINED_MODELS[cls.model_name].keys():
            weights_path = Path(
                hf_hub_download(
                    repo_id=PRETRAINED_MODELS[cls.model_name][weights]["repo_id"],
                    filename=PRETRAINED_MODELS[cls.model_name][weights]["filename"],
                )
            )

        else:
            raise ValueError(
                "Please provide a valid path. or a pre-trained model downloaded from the"
                f" histolytics-hub. One of {list(PRETRAINED_MODELS[cls.model_name].keys())}."
            )

    enc_name, n_nuc_classes, n_tissue_classes, state_dict = cls._get_state_dict(
        weights_path, device=device
    )

    model_inst = cls(
        n_nuc_classes=n_nuc_classes,
        n_tissue_classes=n_tissue_classes,
        enc_name=enc_name,
        enc_pretrain=False,
        enc_freeze=False,
        device=device,
        model_kwargs=model_kwargs,
    )

    if weights_path.suffix == ".safetensors":
        try:
            from safetensors.torch import load_model
        except ImportError:
            raise ImportError(
                "Please install `safetensors` package to load .safetensors files."
            )
        load_model(model_inst.model, weights_path, device.type)
    else:
        model_inst.model.load_state_dict(state_dict, strict=True)

    try:
        cls.nuc_classes = MODEL_CLASS_DICTS[weights]["nuc"]
        cls.tissue_classes = MODEL_CLASS_DICTS[weights]["tissue"]
    except KeyError:
        # if the model is not in the class dict, set to None
        cls.nuc_classes = None
        cls.tissue_classes = None

    return model_inst

predict 

predict(x: Union[Tensor, ndarray, Image], *, use_sliding_win: bool = False, window_size: Tuple[int, int] = None, stride: int = None, save_intermediate: bool = False) -> Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]

Predict the input image or image batch.

Parameters:

Name Type Description Default
x Union[Tensor, ndarray, Image]

Input image (H, W, C) or input image batch (B, C, H, W).

 required
use_sliding_win bool

Whether to use sliding window for prediction.

 False
window_size Tuple[int, int]

The height and width of the sliding window. If use_sliding_win is False this argument is ignored.

 None
stride int

The stride for the sliding window. If use_sliding_win is False this argument is ignored.

 None
save_intermediate bool, default=False

Whether to save intermediate results (logits). If True, the method returns a tuple (final predictions, intermediate results), where the intermediate results are the raw model outputs before argmax.

 False

Returns:

Type Description
Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]

Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]: Dictionary of soft outputs:

- "nuclei": SoftInstanceOutput (type_map, aux_map).
- "tissue": SoftSemanticOutput (type_map).

Examples:

>>> my_model.set_inference_mode()
>>> # with sliding window if image is large
>>> x = my_model.predict(x=image, use_sliding_win=True, window_size=(256, 256), stride=128)
>>> # without sliding window if image is small enough
>>> x = my_model.predict(x=image, use_sliding_win=False)
Source code in src/histolytics/models/_base_model.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def predict(
    self,
    x: Union[torch.Tensor, np.ndarray, Image],
    *,
    use_sliding_win: bool = False,
    window_size: Tuple[int, int] = None,
    stride: int = None,
    save_intermediate: bool = False,
) -> Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]:
    """Predict the input image or image batch.

    Parameters:
        x (Union[torch.Tensor, np.ndarray, Image]):
            Input image (H, W, C) or input image batch (B, C, H, W).
        use_sliding_win (bool):
            Whether to use sliding window for prediction.
        window_size (Tuple[int, int]):
            The height and width of the sliding window. If `use_sliding_win` is False
            this argument is ignored.
        stride (int):
            The stride for the sliding window. If `use_sliding_win` is False this
            argument is ignored.
        save_intermediate (bool, default=False):
            Whether to save intermediate results (logits). If True, the method
            returns a tuple (final predictions, intermediate results), where the
            intermediate results are the raw model outputs before argmax.

    Returns:
        Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]:
            Dictionary of soft outputs:

                - "nuclei": SoftInstanceOutput (type_map, aux_map).
                - "tissue": SoftSemanticOutput (type_map).

    Examples:
        >>> my_model.set_inference_mode()
        >>> # with sliding window if image is large
        >>> x = my_model.predict(x=image, use_sliding_win=True, window_size=(256, 256), stride=128)
        >>> # without sliding window if image is small enough
        >>> x = my_model.predict(x=image, use_sliding_win=False)
    """
    if not self.inference_mode:
        raise ValueError("Run `.set_inference_mode()` before running `predict`")

    if not use_sliding_win:
        x = self.predictor.predict(
            x=x, apply_boundary_weight=False, save_intermediate=save_intermediate
        )
    else:
        if window_size is None:
            raise ValueError(
                "`window_size` must be provided when using sliding window."
            )
        if stride is None:
            raise ValueError("`stride` must be provided when using sliding window.")

        x = self.predictor.predict_sliding_win(
            x=x,
            window_size=window_size,
            stride=stride,
            apply_boundary_weight=True,
            save_intermediate=save_intermediate,
        )

    return x

post_process 

post_process(x: Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]], *, use_async_postproc: bool = True, start_method: str = 'threading', n_jobs: int = 4, save_paths_nuc: List[Union[Path, str]] = None, save_paths_cyto: List[Union[Path, str]] = None, save_paths_tissue: List[Union[Path, str]] = None, coords: List[Tuple[int, int, int, int]] = None, class_dict_nuc: Dict[int, str] = None, class_dict_cyto: Dict[int, str] = None, class_dict_tissue: Dict[int, str] = None, nuc_smooth_func: Callable = gaussian_smooth, cyto_smooth_func: Callable = gaussian_smooth, tissue_smooth_func: Callable = None) -> Dict[str, List[np.ndarray]]

Post-process the output of the model.

Parameters:

Name Type Description Default
x Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]

The output of the .predict() method.

 required
use_async_postproc bool

Whether to use async post-processing. Can give some run-time benefits.

 True
start_method str

The start method. One of: "threading", "fork", "spawn". See mpire docs.

 'threading'
n_jobs int

The number of workers for the post-processing.

 4
save_paths_nuc List[Union[Path, str]]

The paths to save the panlei masks. If None, the masks are not saved.

 None
save_paths_cyto List[Union[Path, str]]

The paths to save the cytoplasm masks. If None, the masks are not saved.

 None
save_paths_tissue List[Union[Path, str]]

The paths to save the tissue masks. If None, the masks are not saved.

 None
coords List[Tuple[int, int, int, int]]

The XYWH coordinates of the image patch. If not None, the coordinates are saved in the filenames of outputs.

 None
class_dict_nuc Dict[int, str]

The dictionary of panlei classes. E.g. {0: "bg", 1: "neoplastic"}

 None
class_dict_cyto Dict[int, str]

The dictionary of cytoplasm classes. E.g. {0: "bg", 1: "macrophage_cyto"}

 None
class_dict_tissue Dict[int, str]

The dictionary of tissue classes. E.g. {0: "bg", 1: "stroma", 2: "tumor"}

 None
nuc_smooth_func Callable

The smoothing function to apply to the nuclei instance maps before post-processing. If None, no smoothing is applied. This is only used when nuclei segmentation masks are saved into vectorized format (e.g. parquet). Ignored save_paths_nuc is None.

 gaussian_smooth
cyto_smooth_func Callable

The smoothing function to apply to the cytoplasm instance maps before post-processing. If None, no smoothing is applied. This is only used when cytoplasm segmentation masks are saved into vectorized format (e.g. parquet). Ignored save_paths_cyto is None.

 gaussian_smooth
tissue_smooth_func Callable

The smoothing function to apply to the tissue type maps before post-processing. If None, no smoothing is applied. This is only used when tissue segmentation masks are saved into vectorized format (e.g. parquet). Ignored save_paths_tissue is None.

 None

Returns:

Type Description
Dict[str, List[ndarray]]

Dict[str, List[np.ndarray]]: Dictionary of post-processed outputs:

  • "nuclei": List of output nuclei masks (H, W).
  • "cyto": List of output cytoplasm masks (H, W).
  • "tissue": List of output tissue masks (H, W).

Examples:

>>> my_model.set_inference_mode()
>>> x = my_model.predict(x=image, use_sliding_win=False)
>>> x = my_model.post_process(
...     x,
...     use_async_postproc=True,
...     start_method="threading",
...     n_jobs=4,
... )
Source code in src/histolytics/models/_base_model.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
def post_process(
    self,
    x: Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]],
    *,
    use_async_postproc: bool = True,
    start_method: str = "threading",
    n_jobs: int = 4,
    save_paths_nuc: List[Union[Path, str]] = None,
    save_paths_cyto: List[Union[Path, str]] = None,
    save_paths_tissue: List[Union[Path, str]] = None,
    coords: List[Tuple[int, int, int, int]] = None,
    class_dict_nuc: Dict[int, str] = None,
    class_dict_cyto: Dict[int, str] = None,
    class_dict_tissue: Dict[int, str] = None,
    nuc_smooth_func: Callable = gaussian_smooth,
    cyto_smooth_func: Callable = gaussian_smooth,
    tissue_smooth_func: Callable = None,
) -> Dict[str, List[np.ndarray]]:
    """Post-process the output of the model.

    Parameters:
        x (Dict[str, Union[SoftSemanticOutput, SoftInstanceOutput]]):
            The output of the .predict() method.
        use_async_postproc (bool):
            Whether to use async post-processing. Can give some run-time benefits.
        start_method (str):
            The start method. One of: "threading", "fork", "spawn". See mpire docs.
        n_jobs (int):
            The number of workers for the post-processing.
        save_paths_nuc (List[Union[Path, str]]):
            The paths to save the panlei masks. If None, the masks are not saved.
        save_paths_cyto (List[Union[Path, str]]):
            The paths to save the cytoplasm masks. If None, the masks are not saved.
        save_paths_tissue (List[Union[Path, str]]):
            The paths to save the tissue masks. If None, the masks are not saved.
        coords (List[Tuple[int, int, int, int]]):
            The XYWH coordinates of the image patch. If not None, the coordinates are
            saved in the filenames of outputs.
        class_dict_nuc (Dict[int, str]):
            The dictionary of panlei classes. E.g. {0: "bg", 1: "neoplastic"}
        class_dict_cyto (Dict[int, str]):
            The dictionary of cytoplasm classes. E.g. {0: "bg", 1: "macrophage_cyto"}
        class_dict_tissue (Dict[int, str]):
            The dictionary of tissue classes. E.g. {0: "bg", 1: "stroma", 2: "tumor"}
        nuc_smooth_func (Callable):
            The smoothing function to apply to the nuclei instance maps before
            post-processing. If None, no smoothing is applied. This is only used when
            nuclei segmentation masks are saved into vectorized format (e.g. parquet).
            Ignored save_paths_nuc is None.
        cyto_smooth_func (Callable):
            The smoothing function to apply to the cytoplasm instance maps before
            post-processing. If None, no smoothing is applied. This is only used when
            cytoplasm segmentation masks are saved into vectorized format (e.g. parquet).
            Ignored save_paths_cyto is None.
        tissue_smooth_func (Callable):
            The smoothing function to apply to the tissue type maps before
            post-processing. If None, no smoothing is applied. This is only used when
            tissue segmentation masks are saved into vectorized format (e.g. parquet).
            Ignored save_paths_tissue is None.

    Returns:
        Dict[str, List[np.ndarray]]:
            Dictionary of post-processed outputs:

            - "nuclei": List of output nuclei masks (H, W).
            - "cyto": List of output cytoplasm masks (H, W).
            - "tissue": List of output tissue masks (H, W).

    Examples:
        >>> my_model.set_inference_mode()
        >>> x = my_model.predict(x=image, use_sliding_win=False)
        >>> x = my_model.post_process(
        ...     x,
        ...     use_async_postproc=True,
        ...     start_method="threading",
        ...     n_jobs=4,
        ... )
    """
    if not self.inference_mode:
        raise ValueError(
            "Run `.set_inference_mode()` before running `post_process`"
        )

    # if batch size is 1, run serially
    if x["tissue"].type_map.shape[0] == 1:
        return self.post_processor.postproc_serial(
            x,
            save_paths_nuc=save_paths_nuc,
            save_paths_cyto=save_paths_cyto,
            save_paths_tissue=save_paths_tissue,
            coords=coords,
            class_dict_nuc=class_dict_nuc,
            class_dict_cyto=class_dict_cyto,
            class_dict_tissue=class_dict_tissue,
        )

    if use_async_postproc:
        x = self.post_processor.postproc_parallel_async(
            x,
            start_method=start_method,
            n_jobs=n_jobs,
            save_paths_nuc=save_paths_nuc,
            save_paths_cyto=save_paths_cyto,
            save_paths_tissue=save_paths_tissue,
            coords=coords,
            class_dict_nuc=class_dict_nuc,
            class_dict_cyto=class_dict_cyto,
            class_dict_tissue=class_dict_tissue,
            nuc_smooth_func=gaussian_smooth,
            cyto_smooth_func=gaussian_smooth,
            tissue_smooth_func=None,
        )
    else:
        x = self.post_processor.postproc_parallel(
            x,
            start_method=start_method,
            n_jobs=n_jobs,
            save_paths_nuc=save_paths_nuc,
            save_paths_cyto=save_paths_cyto,
            save_paths_tissue=save_paths_tissue,
            coords=coords,
            class_dict_nuc=class_dict_nuc,
            class_dict_cyto=class_dict_cyto,
            class_dict_tissue=class_dict_tissue,
            nuc_smooth_func=gaussian_smooth,
            cyto_smooth_func=gaussian_smooth,
            tissue_smooth_func=None,
        )

    return x