o - i>p@sfddlmZddlZddlZddlZddlmZddlmZm Z ddl Z ddl m Z ddl mZddlmZddlmZmZddlmZmZdd lmZmZmZmZdd lmZgd ZGd d d e jZ Gddde j!Z"Gddde j!Z#Gddde j!Z$Gddde jZ%Gddde j&Z'Gddde jZ(d(ddZ)d)d"d#Z*d*d+d&d'Z+e(Z,Z-Z.dS),) annotationsN) OrderedDict)CallableSequence) download_url)UpSample)ConvDropout) get_act_layerget_norm_layer)HoVerNetBranch HoVerNetModeInterpolateMode UpsampleMode)look_up_option)HoVerNetHovernetHoVernetrcs<eZdZddddifdddfdfdd ZdddZZS)_DenseLayerDecoderreluinplaceTbatchr num_featuresint in_channels out_channels dropout_probfloatact str | tuplenorm kernel_sizepaddingreturnNonec stttjdf} ttjdf} t|_|j dt |d|d|j dt |d|j d| ||ddd |j d t |d|d|j d t |d|j d | ||||d dd|dkro|j d| |dSdS)a1 Args: num_features: number of internal channels used for the layer in_channels: number of the input channels. out_channels: number of the output channels. dropout_prob: dropout rate after each dense layer. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. kernel_size: size of the kernel for >1 convolutions (dependent on mode) padding: padding value for >1 convolutions. z preact_bna/bnname spatial_dimschannelszpreact_bna/relur)conv1Fr#biasz conv1/normz conv1/relu2conv2)r#r$groupsr0rdropoutN) super__init__rCONVr DROPOUTnn Sequentiallayers add_moduler r ) selfrrrrr r"r#r$ conv_type dropout_type __class__^/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/hovernet.pyr65s  z_DenseLayerDecoder.__init__x torch.TensorcCsn||}|jd|jdkr-|jd|jdd}|dddd|| || f}t||gd}|S)Nr'r.)r;shapetorchcat)r=rDx1trimrBrBrCforward_s $z_DenseLayerDecoder.forward)rrrrrrrrr r!r"r!r#rr$rr%r&rDrEr%rE__name__ __module__ __qualname__r6rL __classcell__rBrBr@rCr3s *rcs2eZdZddddifdddfdfdd ZZS) _DecoderBlockrrrTrrFr;rrrrrrr r!r"r# same_paddingboolr%r&c  stttjdf} | r|dnd} |d| ||d|| dd|d} t|D]} t|| |||||| d}| |7} |d| d |q+t| ||d }|d ||d | | | d dd dS)a Args: layers: number of layers in the block. num_features: number of internal features used. in_channels: number of the input channel. out_channels: number of the output channel. dropout_prob: dropout rate after each dense layer. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. kernel_size: size of the kernel for >1 convolutions (dependent on mode) same_padding: whether to do padding for >1 convolutions to ensure the output size is the same as the input size. r'rconvar2Fr#r$r0)r r"r#r$zdenselayerdecoder%dr.r r" bna_blockZconvfr/N)r5r6rr7r<ranger _Transition)r=r;rrrrr r"r#rTr>r$ _in_channelsilayertransr@rBrCr6ls.    z_DecoderBlock.__init__)r;rrrrrrrrrr r!r"r!r#rrTrUr%r&rOrPrQr6rRrBrBr@rCrSjs rScs2eZdZddddifdddfdfdd ZZS) _DenseLayerrrrTrrrrrrrrrr r!r"drop_first_norm_relur#r%r&c sZtt|_ttjdf} ttjdf} |s0|j dt |d|d|j dt |d|j d| ||ddd d |j d t |d|d|j d t |d|d krj|rj|j d| |||ddd dn|j d| |||dd d |j dt |d|d|j dt |d|j d| ||ddd d |dkr|j d| |dSdS)a2Dense Convolutional Block. References: Huang, Gao, et al. "Densely connected convolutional networks." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017. Args: num_features: number of internal channels used for the layer in_channels: number of the input channels. out_channels: number of the output channels. dropout_prob: dropout rate after each dense layer. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. drop_first_norm_relu - omits the first norm/relu for the first layer kernel_size: size of the kernel for >1 convolutions (dependent on mode) r'z preact/bnr(z preact/relur,r-r.rFrWzconv1/bnz conv1/relu@r1r#strider$r0zconv2/bnz conv2/reluconv3r4N) r5r6r9r:r;rr7r r8r<r r ) r=rrrrr r"rbr#r>r?r@rBrCr6s.   z_DenseLayer.__init__)rrrrrrrrr r!r"r!rbrr#rr%r&r`rBrBr@rCras racs,eZdZdddifdfdfd d ZZS)r[rrTrrrr r!r"r%r&cs6t|dt|d|d|dt|ddS)z Args: in_channels: number of the input channel. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. bnr'r(rr,N)r5r6r<r r )r=rr r"r@rBrCr6s z_Transition.__init__)rrr r!r"r!r%r&r`rBrBr@rCr[s r[cs<eZdZddddifdddfdfdd ZdddZZS)_ResidualBlockrrrTrFr;rrrrrrr r!r"freeze_dense_layerrU freeze_blockr%r&c stt|_ttjdf} |dkr| ||ddd|_n | ||ddddd|_t||||||dd} |j d | t d|D]} t||||||d } |j d | | qAt |||d |_ |rh|j d| rq| dd Sd S) a>Residual block. References: He, Kaiming, et al. "Deep residual learning for image recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016. Args: layers: number of layers in the block. num_features: number of internal features used. in_channels: number of the input channel. out_channels: number of the output channel. dropout_prob: dropout rate after each dense layer. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. freeze_dense_layer: whether to freeze all dense layers within the block. freeze_block: whether to freeze the whole block. r'rcr.Fr/rdT)r r"rbZ denselayer_0rXZ denselayer_N)r5r6r9r:r;rr7shortcutrar<rZr[rYrequires_grad_) r=r;rrrrr r"rirjr>r^r]r@rBrCr6s&   z_ResidualBlock.__init__rDrEcCs||}|jjdkr|ddddddddf}|jD]+}||}|jdd|jddkrC|ddddddddf}||}|}q||}|S)N)r'r'rF)rkrer;rLrGrY)r=rDscr^rBrBrCrL,s       z_ResidualBlock.forward)r;rrrrrrrrrr r!r"r!rirUrjrUr%r&rMrNrBrBr@rCrhs 9rhcs@eZdZddddifddddd fd fdd Zd!ddZZS)"_DecoderBranch)r2rrTrrr'rF decode_config Sequence[int]r r!r"rrrrr#rTrUr%r&c s@tttjdf}d} d} d} t|_t|D]\} } t| | | | |||||d }|j d| d|d} qt|_ t |}|dd}tt d |d d |dd |d fg}|j d|d|tt dt |dd dfdt|dfd|d |dddfg}|j d|d|tddtjtjd d|_dS)aA Args: decode_config: number of layers for each block. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. dropout_prob: dropout rate after each dense layer. out_channels: number of the output channel. kernel_size: size of the kernel for >1 convolutions (dependent on mode) same_padding: whether to do padding for >1 convolutions to ensure the output size is the same as the input size. r'i ) r;rrrrr r"r#rTZ decoderblockr.irVrcF)r#rer0r$rgr(rr,conv)r#re scale_factormode interp_moder0N)r5r6rr7r9r:decoder_blocks enumeraterSr<output_featureslenrr r rr NONTRAINABLErBILINEARupsample)r=rqr r"rrr#rTr>r\ _num_features _out_channelsr] num_layersblock_i _pad_sizeZ _seq_blockr@rBrCr6AsR        z_DecoderBranch.__init__xinrE short_cutslist[torch.Tensor]cCst|d}|||}|jD]7}||}||}|d8}||jd|jdd}|dkrF|||dddd|| || f7}q|jD]}||}qJ|S)Nr.rFr'r)r~r{rrGr})r=rr block_numberrDrrKrBrBrCrLs    ,  z_DecoderBranch.forward)rqrrr r!r"r!rrrrr#rrTrUr%r&)rrErrr%rErNrBrBr@rCro?s Iroc sXeZdZdZeZeZejddddddifdd d d d d d f d'fd d! Z d(d%d&Z Z S))ra HoVerNet model References: Graham, Simon et al. Hover-net: Simultaneous segmentation and classification of nuclei in multi-tissue histology images, Medical Image Analysis 2019 https://github.com/vqdang/hover_net https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html This network is non-deterministic since it uses `torch.nn.Upsample` with ``UpsampleMode.NONTRAINABLE`` mode which is implemented with torch.nn.functional.interpolate(). Please check the link below for more details: https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms Args: mode: use original implementation (`HoVerNetMODE.ORIGINAL` or "original") or a faster implementation (`HoVerNetMODE.FAST` or "fast"). Defaults to `HoVerNetMODE.FAST`. in_channels: number of the input channel. np_out_channels: number of the output channel of the nucleus prediction branch. out_classes: number of the nuclear type classes. act: activation type and arguments. Defaults to relu. norm: feature normalization type and arguments. Defaults to batch norm. decoder_padding: whether to do padding on convolution layers in the decoders. In the conic branch of the referred repository, the architecture is changed to do padding on convolution layers in order to get the same output size as the input, and this changed version is used on CoNIC challenge. Please note that to get consistent output size, `HoVerNetMode.FAST` mode should be employed. dropout_prob: dropout rate after each dense layer. pretrained_url: if specifying, will loaded the pretrained weights downloaded from the url. There are two supported forms of weights: 1. preact-resnet50 weights coming from the referred hover_net repository, each user is responsible for checking the content of model/datasets and the applicable licenses and determining if suitable for the intended use. please check the following link for more details: https://github.com/vqdang/hover_net#data-format 2. standard resnet50 weights of torchvision. Please check the following link for more details: https://pytorch.org/vision/main/_modules/torchvision/models/resnet.html#ResNet50_Weights adapt_standard_resnet: if the pretrained weights of the encoder follow the original format (preact-resnet50), this value should be `False`. If using the pretrained weights that follow torchvision's standard resnet50 format, this value should be `True`. pretrained_state_dict_key: this arg is used when `pretrained_url` is provided and `adapt_standard_resnet` is True. It is used to extract the expected state dict. freeze_encoder: whether to freeze the encoder of the network. rr'rrrTrFrNryHoVerNetMode | strrrnp_out_channels out_classesr r!r"decoder_paddingrUrrpretrained_url str | Noneadapt_standard_resnetpretrained_state_dict_keyfreeze_encoderr%r&c  stt|tr|}t|t|_|jdkr"|dur"t d|dkr*t d|dkr2t d|dks:|dkr>t d d } d }|jtj krMd }d }nd }d}t t j df}ttd||| dd|ddfdt|d| dfdt|dfg|_| }d}| }t|_t|D]4\}}d}d}| r|dkrd}nd}t|||||||||d }|jd|||}|d9}|d9}qt|_|jd|||dddddtddtjtjdd|_t|||d|_ t||d|_!|dkrt|||dnd|_"|#D]0}t||rtj$%t&'|j(qt|tj)r.tj$*t&'|j(dtj$*t&'|j+dq| durI| r>t,| | d}nt-| }t.||dSdS) NORIGINALTzl'decoder_padding=True' only works when mode is 'FAST', otherwise the output size may not equal to the input.rsz5Number of nuclear types classes exceeds maximum (128)r.z:Number of nuclear type classes should either be None or >1rz+Dropout can only be in the range 0.0 to 1.0rc)rr2rrr'rvFrdrgr(rr,ru) r;rrrrr r"rirjdZconv_bottleneckrw)r#rTr)r#rT)rr#rT)state_dict_key)/r5r6 isinstancestrupperrr rywarningswarn ValueErrorFASTrr7r9r:rr r conv0 res_blocksr|rhr< bottleneckrrrrrrronucleus_predictionhorizontal_verticaltype_predictionmodulesinitkaiming_normal_rH as_tensorweight BatchNorm2d constant_r0_remap_standard_resnet_model_remap_preact_resnet_model_load_pretrained_encoder)r=ryrrrr r"rrrrrrZ_init_featuresZ _block_configZ_ksize_padr>r\rrr]rrirjrmweightsr@rBrCr6s             zHoVerNet.__init__rDrEdict[str, torch.Tensor]cCs|jtjjkr|jddks|jddkrtdn|jddks(|jddkr,td||}g}t|jD]\}}| |}|dkrJ| |q8| |}| |}t jj|||t jj|||i}|jdurv||||t jj<|S)NrFirmz?Input size should be 270 x 270 when using HoVerNetMode.ORIGINALruz;Input size should be 256 x 256 when using HoVerNetMode.FASTr')ryr rvaluerGrrr|rrLappendrrr NPrHVrrNC)r=rDrr]routputrBrBrCrLFs*      zHoVerNet.forward)ryrrrrrrrr r!r"r!rrUrrrrrrUrrrrUr%r&)rDrEr%r) rOrPrQ__doc__r Moder ZBranchrr6rLrRrBrBr@rCrs$+ {rmodel nn.Module state_dictOrderedDict | dictcsr|fddD|tdkr)tddSttdtddS)Ncs2i|]\}}|vr|j|jkr||qSrB)rG).0kv model_dictrrBrC fs,z,_load_pretrained_encoder..rzcno key will be updated. Please check if 'pretrained_url' or `pretrained_state_dict_key` is correct.z out of z* keys are updated with pretrained weights.) ritemsupdateload_state_dictr~keysrrprint)rrrBrrCrds    r model_urlrc Cstd}td}td}td}tjtjd}t|d|ddtj r-dnt d }tj ||dd d }t |D]H}d} ||rSt|d |} n%||rxt|d |} || rlt|d| } n || rxt|d| } | r|||| <||=d|vr||=qB|S)Nz^(conv0\.\/)(.+)$z^(d\d+)\.(.+)$z^(.+\.d\d+)\.units\.(\d+)(.+)$z^(.+\.d\d+)\.blk_bna\.(.+)zpreact-resnet50.pthTFfuzzyfilepathprogresscpu map_location weights_onlydescz conv0.conv\2zres_blocks.\1.\2z \1.layers.denselayer_\2.layers\3z\1.bna_block.\2Z upsample2xrecompileospathjoinrHhubget_dirrcuda is_availabledeviceloadlistrmatchsub) r pattern_conv0 pattern_blockZ pattern_layerZ pattern_bna weights_dirrrkeynew_keyrBrBrCrts2         rrrcCstd}td}td}td}td}td}td}tjtjd} t|d | d d tj rsz._remap_standard_resnet_model..cSs.|dtt|ddd|dS)Nr.r'z.layers.preact/bn.r)rrrrrBrBrCrs.z\1.conv\2/bn.\3z\1.shortcut.\2z\1.bna_block.bn.\2r)rrrZ pattern_bn1rZpattern_block_bn3Zpattern_block_bnZpattern_downsample0Zpattern_downsample1rrrrrrBrBrCrsT               r)rrrr)rr)N)rrrr)/ __future__rrrr collectionsrcollections.abcrrrHtorch.nnr9monai.apps.utilsrmonai.networks.blocksrmonai.networks.layers.factoriesrr monai.networks.layers.utilsr r monai.utils.enumsr r rrmonai.utils.moduler__all__Modulerr:rSrar[rh ModuleListrorrrrrrHoverNetrBrBrBrCs6      79=N] I  6