# -*- coding: utf-8 -*-
# CellViT networks and adaptions
#
# UNETR paper and code: https://github.com/tamasino52/UNETR
# SAM paper and code: https://segment-anything.com/
#
# @ Fabian Hörst, fabian.hoerst@uk-essen.de
# Institute for Artifical Intelligence in Medicine,
# University Medicine Essen
from collections import OrderedDict
from dataclasses import dataclass
from functools import partial
from json import decoder
from pathlib import Path
from typing import List, Literal, Tuple, Union, Optional, Sequence, Callable
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from cell_segmentation.utils.post_proc_cellvit import DetectionCellPostProcessor
from monai.networks.blocks import UpSample
from monai.networks.layers.factories import Conv
from monai.networks.layers.utils import get_act_layer
from monai.networks.nets.basic_unet import UpCat,TwoConv
from monai.utils import InterpolateMode
from .cellvit_unirepLKnet import UniRepLKNet
from .replknet import *
class LayerNorm(nn.Module):
""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
with shape (batch_size, channels, height, width).
"""
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
super().__init__()
self.weight = nn.Parameter(torch.ones(normalized_shape))
self.bias = nn.Parameter(torch.zeros(normalized_shape))
self.eps = eps
self.data_format = data_format
if self.data_format not in ["channels_last", "channels_first"]:
raise NotImplementedError
self.normalized_shape = (normalized_shape, )
def forward(self, x):
if self.data_format == "channels_last":
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
elif self.data_format == "channels_first":
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x
encoder_feature_channel = {
"unireplknet_a": (40, 80, 160, 320),
"unireplknet_f": (48, 96, 192, 384),
"unireplknet_p": (64, 128, 256, 512),
"unireplknet_n": (80, 160, 320, 640),
"unireplknet_t": (80, 160, 320, 640),
"unireplknet_s": (96, 192, 384, 768),
"unireplknet_b": (128, 256, 512, 1024),
"unireplknet_l": (192, 384, 768, 1536),
"unireplknet_xl": (256, 512, 1024, 2048),
"convnext_tiny": (48, 96, 192, 384, 768),
"resnet50": (64, 256, 512, 1024, 2048),
"vitdet_small": (768, 768, 768, 768, 768),
}
def _get_encoder_channels_by_backbone(backbone: str, in_channels: int = 3) -> tuple:
"""
Get the encoder output channels by given backbone name.
Args:
backbone: name of backbone to generate features, can be from [efficientnet-b0, ..., efficientnet-b7].
in_channels: channel of input tensor, default to 3.
Returns:
A tuple of output feature map channels' length .
"""
encoder_channel_tuple = encoder_feature_channel[backbone] #encoder_channel_tuple是指的编码器通道元组,在这里是有4个通道的元组
encoder_channel_list = [in_channels] + list(encoder_channel_tuple) #encoder_channel_list是指的编码器通道列表[3,80,160,320,640]
encoder_channel = tuple(encoder_channel_list) #encoder_channel是指的编码器通道元组(3,80,160,320,640)
return encoder_channel
class RepLKDeocder(nn.Module):
def __init__(self,
encoder_channels: Sequence[int],
spatial_dims: int,
decoder_channels: Sequence[int],
stage_lk_sizes,
drop_path,
upsample: str,
pre_conv: Optional[str],
interp_mode: str,
align_corners: Optional[bool],
small_kernel,
dw_ratio: int=1,
small_kernel_merged=False,
norm: Union[str, tuple] = ("batch", {"eps": 1e-3, "momentum": 0.1}),
act: Union[str, tuple] = ("relu", {"inplace": True}),
dropout: Union[float, tuple] = 0.0,
bias: bool = False,
is_pad: bool = True,
ffn_ratio=4,
):
super().__init__()
in_channels = [encoder_channels[-1]] + list(decoder_channels[:-1]) #in_channels=[640,1024,512,256,128]
skip_channels = list(encoder_channels[1:-1][::-1]) + [0]
halves = [True] * (len(skip_channels) - 1)
halves.append(False)
stage_lk_sizes = stage_lk_sizes
blocks = []
for in_chn, skip_chn, out_chn, halve in zip(in_channels, skip_channels, decoder_channels, halves):
blocks.append(
UpCat(
spatial_dims=spatial_dims,
in_chns=in_chn,
cat_chns=skip_chn,
out_chns=out_chn,
act=act,
norm=norm,
dropout=dropout,
bias=bias,
upsample=upsample,
pre_conv=pre_conv,
interp_mode=interp_mode,
align_corners=align_corners,
halves=halve,
is_pad=is_pad,
)
)
self.blocks = nn.ModuleList(blocks)
repblock = []
for i in range(4):
repblock.append(RepLKBlock(in_channels=in_channels[i], dw_channels=int(in_channels[i] * dw_ratio), block_lk_size=stage_lk_sizes[i],
small_kernel=small_kernel, drop_path=drop_path, small_kernel_merged=small_kernel_merged))
self.repblock = nn.ModuleList(repblock)
convffnblock = []
for i in range(4):
convffnblock.append(ConvFFN(in_channels=in_channels[i], internal_channels=int(in_channels[i] * ffn_ratio), out_channels=in_channels[i], drop_path=drop_path))
self.convffnblock = nn.ModuleList(convffnblock)
self.upsample = [UpSample(
spatial_dims,
decoder_channels[i],
decoder_channels[i],
mode=upsample,
pre_conv=pre_conv,
interp_mode=interp_mode,
align_corners=align_corners,) for i in range(len(decoder_channels) - 1)]
self.upsample1= UpSample(
spatial_dims,
256,
256,
2,
mode=upsample,
pre_conv=pre_conv,
interp_mode=interp_mode,
align_corners=align_corners,
)
self.upsample2= UpSample(
spatial_dims,
128,
128,
2,
mode=upsample,
pre_conv=pre_conv,
interp_mode=interp_mode,
align_corners=align_corners,
)
self.convs = TwoConv(spatial_dims, 304, decoder_channels[-2], act, norm, bias, dropout)
self.convs1 = TwoConv(spatial_dims, 152, decoder_channels[-1], act, norm,bias, dropout)
def forward(self, features: List[torch.Tensor], input_feature: torch.Tensor, skip_connect: int = 3):
skips = features[:-1][::-1] #skips[0],[1],[2]=[16,320,16,16],[16,160,32,32],[16,80,64,64],[16,40,128,128]
features = features[1:][::-1]
#input_feature = self.conv1(input_feature) #input_feature=[16,64,256,256]
x = features[0] #x=[16,640,8,8]
for i, (block, repblock, convffnblock) in enumerate(zip(self.blocks, self.repblock, self.convffnblock)):
if i < skip_connect:
skip = skips[i] #skip=[16,320,16,16], skip=[16,160,32,32], skip=[16,80,64,64], skip = [16,40,128,128]
#x = repblock(x)
#x = convffnblock(x)
x = block(x, skip)
else:
#x = repblock(x)
skip = input_feature[1]
x =self.upsample1(x)
x = torch.cat([skip, x], dim=1)
x = self.convs(x)
skip = input_feature[0]
x = self.upsample2(x)
x = torch.cat([skip, x], dim=1)
x = self.convs1(x)
return x
class SegmentationHead(nn.Sequential):
"""
Segmentation head.
This class refers to `segmentation_models.pytorch
<https://github.com/qubvel/segmentation_models.pytorch>`_.
Args:
spatial_dims: number of spatial dimensions.
in_channels: number of input channels for the block.
out_channels: number of output channels for the block.
kernel_size: kernel size for the conv layer.
act: activation type and arguments.
scale_factor: multiplier for spatial size. Has to match input size if it is a tuple.
"""
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
act: Optional[Union[Tuple, str]] = None,
scale_factor: float = 1.0,
):
conv_layer = Conv[Conv.CONV, spatial_dims](
in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=kernel_size // 2
)
bn_layer = nn.BatchNorm2d(in_channels)
conv_layer1 = conv_bn(in_channels=in_channels, out_channels=in_channels, kernel_size=1, stride=1, padding=0, groups=1)
nonlinear_layer = nn.GELU()
conv_layer2 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, groups=1)
up_layer: nn.Module = nn.Identity()
if scale_factor > 1.0:
up_layer = UpSample(
spatial_dims=spatial_dims,
scale_factor=scale_factor,
mode="nontrainable",
pre_conv=None,
interp_mode=InterpolateMode.LINEAR,
)
if act is not None:
act_layer = get_act_layer(act)
else:
act_layer = nn.Identity()
super().__init__(bn_layer, conv_layer1, nonlinear_layer, conv_layer2, up_layer)
class CellViT(UniRepLKNet):
def __init__(
self,
model256_path: Union[Path, str],
num_nuclei_classes=int,
num_tissue_classes=int,
in_channels=int,
backbone = "unireplknet_s",
decoder_channels: Tuple = (1024, 512, 256, 128, 64),
spatial_dims: int = 2,
norm: Union[str, tuple] = ("batch", {"eps": 1e-3, "momentum": 0.1}),
act: Union[str, tuple] = ("relu", {"inplace": True}),
dropout: Union[float, tuple] = 0.0,
decoder_bias: bool = False,
upsample: str = "nontrainable",
interp_mode: str = "nearest",
drop_path_rate: float = 0.1,
large_kernel_sizes=[13,27,29,31],
small_kernel=5,
):
super().__init__()
self.num_tissue_classes = num_tissue_classes
self.num_nuclei_classes = num_nuclei_classes
if backbone not in encoder_feature_channel:
raise ValueError(f"invalid model_name {backbone} found, must be one of {encoder_feature_channel.keys()}.")
if spatial_dims not in (2, 3):
raise ValueError("spatial_dims can only be 2 or 3.")
self.backbone = backbone
self.spatial_dims = spatial_dims
encoder_channels = _get_encoder_channels_by_backbone(backbone, in_channels)
self.encoder = UniRepLKNet(
in_chans=3,
num_classes=None,
depths=(3, 3, 27, 3),
dims=(96,192,384,768),
drop_path_rate=0.3,
layer_scale_init_value=1e-6,
head_init_scale=1.,
kernel_sizes=None,
deploy=False,
with_cp=False,
init_cfg=None,
attempt_use_lk_impl=True,
use_sync_bn=False,
)
self.decoder = RepLKDeocder(
encoder_channels=encoder_channels,
stage_lk_sizes=large_kernel_sizes,
small_kernel=small_kernel,
drop_path=drop_path_rate,
spatial_dims=spatial_dims,
decoder_channels=decoder_channels,
act=act,
norm=norm,
dropout=dropout,
bias=decoder_bias,
upsample=upsample,
interp_mode=interp_mode,
pre_conv=None,
align_corners=None,
)
self.branches_output = {
"nuclei_binary_map": 2,
"hv_map": 2,
"nuclei_type_maps": self.num_nuclei_classes,
} # number of output channels for each branch
self.nuclei_binary_segmentation_head = SegmentationHead(
spatial_dims=spatial_dims,
in_channels=decoder_channels[-1],
out_channels=self.branches_output["nuclei_binary_map"],
kernel_size=1,
act=None,
scale_factor=1.0,
)
self.hv_map_head = SegmentationHead(
spatial_dims=spatial_dims,
in_channels=decoder_channels[-1],
out_channels=self.branches_output["hv_map"],
kernel_size=1,
act=None,
scale_factor=1.0,
)
self.nuclei_type_maps_head = SegmentationHead(
spatial_dims=spatial_dims,
in_channels=decoder_channels[-1],
out_channels=self.branches_output["nuclei_type_maps"],
kernel_size=1,
act=None,
scale_factor=1.0,
)
self.model256_path = model256_path
def forward(self, x: torch.Tensor) -> dict:
"""Forward pass of CellViT
Args:
x (torch.Tensor): Images in BCHW style
retrieve_tokens (bool, optional): If tokens of ViT should be returned as well. Defaults to False.
Returns:
dict: Output for all branches:
* tissue_types: Raw tissue type prediction. Shape: (batch_size, num_tissue_classes)
* nuclei_binary_map: Raw binary cell segmentation predictions. Shape: (batch_size, 2, H, W)
* hv_map: Binary HV Map predictions. Shape: (batch_size, 2, H, W)
* nuclei_type_map: Raw binary nuclei type preditcions. Shape: (batch_size, num_nuclei_classes, H, W)
* (optinal) tokens
"""
out_dict = {}
#x = x.to(torch.bfloat16)
#self.encoder = self.encoder.to(torch.bfloat16)
classifier_logits, z, input_feature = self.encoder(x)
#classifier_logits = classifier_logits.to(torch.bfloat16)
#z = [z.to(torch.bfloat16) for z in z]
#input_feature = [input_feature.to(torch.bfloat16) for input_feature in input_feature]
# if torch.any(torch.tensor([torch.any(torch.isnan(t)) for t in z])):
# print("z is nan")
out_dict["tissue_types"] = classifier_logits
decoder_output = self.decoder(z,input_feature)
#decoder_output = decoder_output.to(torch.bfloat16)
# if torch.isnan(decoder_output).any():
# print("decoder_output is nan")
out_dict["nuclei_binary_map"] = self.nuclei_binary_segmentation_head(decoder_output)
out_dict["hv_map"] = self.hv_map_head(decoder_output)
out_dict["nuclei_type_map"] = self.nuclei_type_maps_head(decoder_output)
return out_dict
def load_pretrained_encoder(self, model256_path: str):
"""Load pretrained ViT-256 from provided path
Args:
model256_path (str): Path to ViT-256
"""
state_dict = torch.load(str(model256_path), map_location="cpu")
state_dict = {key: value for key, value in state_dict.items() if not key.startswith("head")}
msg = self.encoder.load_state_dict(state_dict, strict=False)
print(f"Loading checkpoint: {msg}")
def reshape_model_output(
self,
predictions: OrderedDict,
device: str,
) -> OrderedDict:
predictions = OrderedDict(
[
[k, v.permute(0, 2, 3, 1).contiguous().to(device)]
for k, v in predictions.items()
if k != "tissue_types"
]
)
return predictions
def calculate_instance_map(
self, predictions: OrderedDict, magnification: Literal[20, 40] = 40
) -> Tuple[torch.Tensor, List[dict]]:
"""Calculate Instance Map from network predictions (after Softmax output)
Args:
predictions (dict): Dictionary with the following required keys:
* nuclei_binary_map: Binary Nucleus Predictions. Shape: (batch_size, H, W, 2)
* nuclei_type_map: Type prediction of nuclei. Shape: (batch_size, H, W, 6)
* hv_map: Horizontal-Vertical nuclei mapping. Shape: (batch_size, H, W, 2)
magnification (Literal[20, 40], optional): Which magnification the data has. Defaults to 40.
Returns:
Tuple[torch.Tensor, List[dict]]:
* torch.Tensor: Instance map. Each Instance has own integer. Shape: (batch_size, H, W)
* List of dictionaries. Each List entry is one image. Each dict contains another dict for each detected nucleus.
For each nucleus, the following information are returned: "bbox", "centroid", "contour", "type_prob", "type"
"""
# reshape to B, H, W, C
predictions_ = predictions.copy()
predictions_["nuclei_type_map"] = predictions_["nuclei_type_map"].permute(
0, 2, 3, 1
)
predictions_["nuclei_binary_map"] = predictions_["nuclei_binary_map"].permute(
0, 2, 3, 1
)
predictions_["hv_map"] = predictions_["hv_map"].permute(0, 2, 3, 1)
predictions = predictions_
cell_post_processor = DetectionCellPostProcessor(
nr_types=self.num_nuclei_classes, magnification=magnification, gt=False
)
instance_preds = []
type_preds = []
for i in range(predictions["nuclei_binary_map"].shape[0]):
pred_map = np.concatenate(
[
torch.argmax(predictions["nuclei_type_map"], dim=-1)[i]
.detach()
.cpu()[..., None],
torch.argmax(predictions["nuclei_binary_map"], dim=-1)[i]
.detach()
.cpu()[..., None],
predictions["hv_map"][i].detach().cpu(),
],
axis=-1,
)
instance_pred = cell_post_processor.post_process_cell_segmentation(pred_map)
instance_preds.append(instance_pred[0])
type_preds.append(instance_pred[1])
return torch.Tensor(np.stack(instance_preds)), type_preds
def generate_instance_nuclei_map(
self, instance_maps: torch.Tensor, type_preds: List[dict]
) -> torch.Tensor:
"""Convert instance map (binary) to nuclei type instance map
Args:
instance_maps (torch.Tensor): Binary instance map, each instance has own integer. Shape: (batch_size, H, W)
type_preds (List[dict]): List (len=batch_size) of dictionary with instance type information (compare post_process_hovernet function for more details)
Returns:
torch.Tensor: Nuclei type instance map. Shape: (batch_size, H, W, self.num_nuclei_classes)
"""
batch_size, h, w = instance_maps.shape
instance_type_nuclei_maps = torch.zeros(
(batch_size, h, w, self.num_nuclei_classes)
)
for i in range(batch_size):
instance_type_nuclei_map = torch.zeros((h, w, self.num_nuclei_classes))
instance_map = instance_maps[i]
type_pred = type_preds[i]
for nuclei, spec in type_pred.items():
nuclei_type = spec["type"]
instance_type_nuclei_map[:, :, nuclei_type][
instance_map == nuclei
] = nuclei
instance_type_nuclei_maps[i, :, :, :] = instance_type_nuclei_map
return instance_type_nuclei_maps
def freeze_encoder(self):
"""Freeze encoder to not train it """
for layer_name, p in self.encoder.named_parameters():
if layer_name.split(".")[0] != "head":
p.requires_grad = False
def unfreeze_encoder(self):
"""Unfreeze encoder to train the whole model """
for p in self.encoder.parameters():
p.requires_grad = True
@dataclass
class DataclassHVStorage:
"""Storing PanNuke Prediction/GT objects for calculating loss, metrics etc. with HoverNet networks
Args:
nuclei_binary_map (torch.Tensor): Softmax output for binary nuclei branch. Shape: (batch_size, 2, H, W)
hv_map (torch.Tensor): Logit output for HV-Map. Shape: (batch_size, 2, H, W)
nuclei_type_map (torch.Tensor): Softmax output for nuclei type-prediction. Shape: (batch_size, num_tissue_classes, H, W)
tissue_types (torch.Tensor): Logit tissue prediction output. Shape: (batch_size, num_tissue_classes)
instance_map (torch.Tensor): Pixel-wise nuclear instance segmentation.
Each instance has its own integer, starting from 1. Shape: (batch_size, H, W)
instance_types_nuclei: Pixel-wise nuclear instance segmentation predictions, for each nuclei type.
Each instance has its own integer, starting from 1.
Shape: (batch_size, num_nuclei_classes, H, W)
batch_size (int): Batch size of the experiment
instance_types (list, optional): Instance type prediction list.
Each list entry stands for one image. Each list entry is a dictionary with the following structure:
Main Key is the nuclei instance number (int), with a dict as value.
For each instance, the dictionary contains the keys: bbox (bounding box), centroid (centroid coordinates),
contour, type_prob (probability), type (nuclei type)
Defaults to None.
regression_map (torch.Tensor, optional): Regression map for binary prediction map.
Shape: (batch_size, 2, H, W). Defaults to None.
regression_loss (bool, optional): Indicating if regression map is present. Defaults to False.
h (int, optional): Height of used input images. Defaults to 256.
w (int, optional): Width of used input images. Defaults to 256.
num_tissue_classes (int, optional): Number of tissue classes in the data. Defaults to 19.
num_nuclei_classes (int, optional): Number of nuclei types in the data (including background). Defaults to 6.
"""
nuclei_binary_map: torch.Tensor
hv_map: torch.Tensor
tissue_types: torch.Tensor
nuclei_type_map: torch.Tensor
instance_map: torch.Tensor
instance_types_nuclei: torch.Tensor
batch_size: int
instance_types: list = None
regression_map: torch.Tensor = None
regression_loss: bool = False
h: int = 256
w: int = 256
num_tissue_classes: int = 19
num_nuclei_classes: int = 6
def get_dict(self) -> dict:
"""Return dictionary of entries"""
property_dict = self.__dict__
if not self.regression_loss and "regression_map" in property_dict.keys():
property_dict.pop("regression_map")
return property_dict