U Ph p@sddlmZddlZddlZddlZddlmZddlmZm Z ddl Z ddl m Z ddl mZddlmZddlmZmZddlmZmZdd lmZmZmZmZdd lmZmZd d d d gZGddde j Z!Gddde j"Z#Gddde j"Z$Gddde j"Z%Gddde j Z&Gddde j'Z(edGdd d e j Z)ddddd Z*d!d"d#d$Z+d)d!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)exportlook_up_optionHoVerNetHovernetHoVernetc sVeZdZddddifdddfdddd d d ddd d fd d ZdddddZZS)_DenseLayerDecoderreluinplaceTbatchrintfloat str | tupleNone) num_features in_channels out_channels dropout_probactnorm kernel_sizepaddingreturnc 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|dkr|j d| |dS)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&groupsr1rdropoutN) super__init__rCONVr DROPOUTnn Sequentiallayers add_moduler r ) selfrr r!r"r#r$r%r& conv_type dropout_type __class__Q/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/hovernet.pyr75s  z_DenseLayerDecoder.__init__ torch.Tensorxr'cCsn||}|jd|jdkrZ|jd|jdd}|dddd|| || f}t||gd}|S)Nr(r/)r<shapetorchcat)r>rGx1trimrCrCrDforward_s  $z_DenseLayerDecoder.forward__name__ __module__ __qualname__r7rN __classcell__rCrCrArDr3s $*rc sHeZdZddddifdddfddddd d d dd d d fdd ZZS) _DecoderBlockrrrTrrFrrrboolr) r<rr r!r"r#r$r% same_paddingr'c  stttjdf} | r$|dnd} |d| ||d|| dd|d} t|D]8} t|| |||||| d}| |7} |d| d |qVt| ||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(rconvar3Fr%r&r1)r#r$r%r&zdenselayerdecoder%dr/r#r$ bna_blockZconvfr0N)r6r7rr8r=ranger _Transition)r>r<rr r!r"r#r$r%rVr?r& _in_channelsilayertransrArCrDr7ls0    z_DecoderBlock.__init__rPrQrRr7rSrCrCrArDrTjs  rTc sFeZdZddddifdddfdddd d d ddd d fd d ZZS) _DenseLayerrrrTrrrrrrr) rr r!r"r#r$drop_first_norm_relur%r'c sXtt|_ttjdf} ttjdf} |s`|j dt |d|d|j dt |d|j d| ||ddd d |j d t |d|d|j d t |d|d kr|r|j d| |||ddd dn|j d| |||dd d |j dt |d|d|j dt |d|j d| ||ddd d |dkrT|j d| |dS)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/rFrXzconv1/bnz conv1/relu@r2r%strider&r1zconv2/bnz conv2/reluconv3r5N) r6r7r:r;r<rr8r r9r=r r ) r>rr r!r"r#r$rcr%r?r@rArCrDr7s0    z_DenseLayer.__init__rarCrCrArDrbs  rbcs6eZdZdddifdfdddddfd 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)r6r7r=r r )r>r r#r$rArCrDr7s z_Transition.__init__rarCrCrArDr\s r\c sXeZdZddddifdddfdddddd d d d d d fd d ZdddddZZS)_ResidualBlockrrrTrFrrrrUr) r<rr r!r"r#r$freeze_dense_layer freeze_blockr'c stt|_ttjdf} |dkr>| ||ddd|_n| ||ddddd|_t||||||dd} |j d | t d|D],} t||||||d } |j d | | qt |||d |_ |r|j d| r| dd 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(rdr/Fr0reT)r#r$rcZ denselayer_0rYZ denselayer_N)r6r7r:r;r<rr8shortcutrbr=r[r\rZrequires_grad_) r>r<rr r!r"r#r$rjrkr?r_r^rArCrDr7s0   z_ResidualBlock.__init__rErFcCs||}|jjdkr6|ddddddddf}|jD]V}||}|jdd|jddkr|ddddddddf}||}|}q<||}|S)N)r(r(rH)rlrfr<rNrIrZ)r>rGscr_rCrCrDrN,s       z_ResidualBlock.forwardrOrCrCrArDris &9ric sZeZdZddddifddddd fd d d d d d dddfdd ZddddddZZS)_DecoderBranch)r3rrTrrr(rFz Sequence[int]rrrrUr) decode_configr#r$r"r!r%rVr'c s@tttjdf}d} d} d} t|_t|D]>\} } t| | | | |||||d }|j d| d|d} q6t|_ 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<rr r!r"r#r$r%rVZ decoderblockr/irWrdF)r%rfr1r&rhr)rr-conv)r%rf scale_factormode interp_moder1N)r6r7rr8r:r;decoder_blocks enumeraterTr=output_featureslenrr r rr NONTRAINABLErBILINEARupsample)r>rrr#r$r"r!r%rVr?r] _num_features _out_channelsr^ num_layersblock_i _pad_sizeZ _seq_blockrArCrDr7AsZ       z_DecoderBranch.__init__rEzlist[torch.Tensor])xin short_cutsr'cCst|d}|||}|jD]n}||}||}|d8}||jd|jdd}|dkr|||dddd|| || f7}q|jD] }||}q|S)Nr/rHr(r)r~r{rrIr})r>rr block_numberrGrrMrCrCrDrNs    .  z_DecoderBranch.forwardrOrCrCrArDrp?s "Irpzmonai.networks.netscszeZdZdZeZeZejddddddifdd d d d d d f d d d d dddddddddd fdd Z dddddZ 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(rrrTrFrNzHoVerNetMode | strrrrUr str | Noner) ryr np_out_channels out_classesr#r$decoder_paddingr"pretrained_urladapt_standard_resnetpretrained_state_dict_keyfreeze_encoderr'c  stt|tr|}t|t|_|jdkrD|dkrDt d|dkrVt dn|dkrft d|dksv|dkr~t d d } d }|jtj krd }d }nd }d}t t j df}ttd||| dd|ddfdt|d| dfdt|dfg|_| }d}| }t|_t|D]n\}}d}d}| r<|dkr8d}nd}t|||||||||d }|jd|||}|d9}|d9}qt|_|jd|||dddddtddtjtjdd|_t|||d|_ t||d|_!|dkrt|||dnd|_"|#D]b}t||r,tj$%t&'|j(n:t|tj)rtj$*t&'|j(dtj$*t&'|j+dq| dk r| rt,| | d}nt-| }t.||dS) 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.0rd)rr3rrr(rvFrerhr)rr-ru) r<rr r!r"r#r$rjrkdZconv_bottleneckrw)r%rVr!)r%rV)r!r%rV)state_dict_key)/r6r7 isinstancestrupperrr rywarningswarn ValueErrorFASTrr8r:r;rr r conv0 res_blocksr|rir= bottleneckrrrrrrrpnucleus_predictionhorizontal_verticaltype_predictionmodulesinitkaiming_normal_rJ as_tensorweight BatchNorm2d constant_r1_remap_standard_resnet_model_remap_preact_resnet_model_load_pretrained_encoder)r>ryr rrr#r$rr"rrrrZ_init_featuresZ _block_configZ_ksize_padr?r]rrr^rrjrkrmweightsrArCrDr7s               zHoVerNet.__init__rEzdict[str, torch.Tensor]rFcCs|jtjjkr4|jddks*|jddkrXtdn$|jddksP|jddkrXtd||}g}t|jD]$\}}| |}|dkrp| |qp| |}| |}t jj|||t jj|||i}|jdk r||||t jj<|S)NrHirnz?Input size should be 270 x 270 when using HoVerNetMode.ORIGINALruz;Input size should be 256 x 256 when using HoVerNetMode.FASTr()ryr rvaluerIrrr|rrNappendrrr NPrHVrrNC)r>rGrr^routputrCrCrDrNGs*         zHoVerNet.forward) rPrQrR__doc__r Moder ZBranchrr7rNrSrCrCrArDrs"+ ,{z nn.ModulezOrderedDict | dictmodel state_dictcsp|fddD|tdkrPtdnttdtddS)Ncs2i|]*\}}|kr|j|jkr||qSrC)rI).0kv model_dictrrCrD gs 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~keysrrprintrrCrrDres   rr) model_urlc Cstd}td}td}td}tjtjd}t|d|ddtj |tj r`dnt d d d }t |D]}d}||rt|d |}nJ||rt|d |}||rt|d|}n||rt|d|}|r||||<||=d|kr~||=q~|S)Nz^(conv0\.\/)(.+)$z^(d\d+)\.(.+)$z^(.+\.d\d+)\.units\.(\d+)(.+)$z^(.+\.d\d+)\.blk_bna\.(.+)zpreact-resnet50.pthTFfuzzyfilepathprogresscpu map_locationdescz conv0.conv\2zres_blocks.\1.\2z \1.layers.denselayer_\2.layers\3z\1.bna_block.\2Z upsample2xrecompileospathjoinrJhubget_dirrloadcuda is_availabledevicelistrmatchsub) r pattern_conv0 pattern_blockZ pattern_layerZ pattern_bna weights_dirrkeynew_keyrCrCrDrus2          rr)rrc Cstd}td}td}td}td}td}td}tjtjd} t|d | d d tj | tj r~dnt d d } |dk r| |} t | D]} d} || rt|d| } n|| rt|d| } n|| rtt|dd| } || r"t|dd| } nR|| r>t|d| } n6|| rZt|d| } n|| rtt|d| } | r| | | | <| | =q| S)Nz ^conv1\.(.+)$z ^bn1\.(.+)$z^layer(\d+)\.(\d+)\.(.+)$z@^(res_blocks.d\d+\.layers\.denselayer_)(\d+)\.layers\.bn3\.(.+)$zB^(res_blocks.d\d+\.layers\.denselayer_\d+\.layers)\.bn(\d+)\.(.+)$z)^(res_blocks.d\d+).+\.downsample\.0\.(.+)z)^(res_blocks.d\d+).+\.downsample\.1\.(.+)z resnet50.pthTFrrrz conv0.conv.\1z conv0.bn.\1cSs6dtt|ddd|dd|dS)Nz res_blocks.dr/z.layers.denselayer_r(z.layers.r)rrgroupsrCrCrDsz._remap_standard_resnet_model..cSs.|dtt|ddd|dS)Nr/r(z.layers.preact/bn.r)rrrrrCrCrDrz\1.conv\2/bn.\3z\1.shortcut.\2z\1.bna_block.bn.\2r) rrrZ pattern_bn1rZpattern_block_bn3Zpattern_block_bnZpattern_downsample0Zpattern_downsample1rrrrrCrCrDrsP       "        r)N)0 __future__rrrr collectionsrcollections.abcrrrJtorch.nnr:monai.apps.utilsrZmonai.networks.blocksrmonai.networks.layers.factoriesrr monai.networks.layers.utilsr r monai.utils.enumsr r rrmonai.utils.modulerr__all__Modulerr;rTrbr\ri ModuleListrprrrrrrHoverNetrCrCrCrDs6      79=N]I5