U Ph>l @sddlmZddlZddlZddlmZddlmZddlm Z ddl m Z ddl Z ddl mZddlmZddlmZmZmZdd lmZmZdd lmZdd lmZmZed d d\ZZeddd\ZZdZ dZ!ddddddddddg Z"dddddgdd d!fdd"d"d"d"gd#d!d!fdd$d%d&d$gd#d!d!fd'd$d%d&d$gdd d!fd'd$d%d(d$gdd d fd'd$d)d*d$gdd d fd'd$d+d*d$gdd d fd,Z#e$e%Z&d-d.Z'd/d0Z(Gd1ddej)Z*Gd2ddej)Z+Gd3ddej)Z,Gd4d5d5e,Z-Gd6d7d7e-eZ.d8d9d:d:d;dd?d@Z/dVdr?r@rA conv_type norm_type __class__r1r2rGJs    zResNetBlock.__init__ torch.TensorxrBcCs`|}||}||}||}||}||}|jdk rJ||}||7}||}|SN)rJrKrArLrMr@rOrVresidualoutr1r1r2forwardis        zResNetBlock.forward__name__ __module__ __qualname__ expansionrGr[ __classcell__r1r1rRr2rGs c sTeZdZdZddddddiffddddd d d d fd d ZdddddZZS)rr&r%r!Nr5r6Tr7r8r9r:r;c stttj|f}ttj|f}|||ddd|_|||_|||d|ddd|_|||_ ||||j ddd|_ |||j |_ t |d|_||_||_dS)a Args: in_planes: number of input channels. planes: number of output channels (taking expansion into account). spatial_dims: number of spatial dimensions of the input image. stride: stride to use for second conv layer. downsample: which downsample layer to use. act: activation type and arguments. Defaults to relu. r!F)rCrEr%rCr?rDrErN)rFrGrrHr rIrJrKrLrMr`conv3bn3r rAr@r?rNrRr1r2rGs    zResNetBottleneck.__init__rTrUcCs~|}||}||}||}||}||}||}||}||}|jdk rh||}||7}||}|SrW)rJrKrArLrMrcrdr@rXr1r1r2r[s           zResNetBottleneck.forwardr\r1r1rRr2r|s "cseZdZdZddddddddd d d d d iff d d d dddddddddddddfdd Zd$ddddddddZd%dddddddddd Zddd!d"d#ZZS)&rag ResNet based on: `Deep Residual Learning for Image Recognition `_ and `Can Spatiotemporal 3D CNNs Retrace the History of 2D CNNs and ImageNet? `_. Adapted from ``_. Args: block: which ResNet block to use, either Basic or Bottleneck. ResNet block class or str. for Basic: ResNetBlock or 'basic' for Bottleneck: ResNetBottleneck or 'bottleneck' layers: how many layers to use. block_inplanes: determine the size of planes at each step. Also tunable with widen_factor. spatial_dims: number of spatial dimensions of the input image. n_input_channels: number of input channels for first convolutional layer. conv1_t_size: size of first convolution layer, determines kernel and padding. conv1_t_stride: stride of first convolution layer. no_max_pool: bool argument to determine if to use maxpool layer. shortcut_type: which downsample block to use. Options are 'A', 'B', default to 'B'. - 'A': using `self._downsample_basic_block`. - 'B': kernel_size 1 conv + norm. widen_factor: widen output for each layer. num_classes: number of output (classifications). feed_forward: whether to add the FC layer for the output, default to `True`. bias_downsample: whether to use bias term in the downsampling block when `shortcut_type` is 'B', default to `True`. act: activation type and arguments. Defaults to relu. r%r!Fr"g?iTr5r6z*type[ResNetBlock | ResNetBottleneck] | str list[int]r7ztuple[int] | intboolstrfloatr9r:)blocklayersblock_inplanesr>n_input_channels conv1_t_sizeconv1_t_stride no_max_pool shortcut_type widen_factor num_classes feed_forwardbias_downsamplerArBcsJtt|tr<|dkr"t}n|dkr0t}n td|ttj|f}t t j |f}t t j |f}t t j |f}t}fdd|D}|d|_||_| |_t||}t||}|||j||tdd|Dd d |_||j|_t|d |_|d d dd|_|||d|d|| |_|j||d|d|| d d|_|j||d |d || d d|_|j||d |d || d d|_||||_| rt |d |j!| nd|_"|#D]}t||rtj$j%t&'|j(dddn^t||rtj$)t&'|j(dtj$)t&'|j*dn$t|tj rtj$)t&'|j*dqdS)Nr r(z+Unknown block '%s', use basic or bottleneckcsg|]}t|qSr1)r7).0rVrrr1r2 sz#ResNet.__init__..rcss|]}|dVqdS)r#Nr1)rvkr1r1r2 sz"ResNet.__init__..Frbrr%r#r!)rCr?rD)r?fan_outr5)mode nonlinearity)+rFrG isinstancerhrr ValueErrorrrHr rIr MAXZ ADAPTIVEAVGr4r<rprur tuplerJrKr rAmaxpool _make_layerlayer1layer2layer3layer4avgpoolnnLinearr`fcmodulesinitkaiming_normal_torch as_tensorweight constant_rE)rOrjrkrlr>rmrnrorprqrrrsrtrurArPrQ pool_typeZ avgp_typeZ block_avgpoolZconv1_kernel_sizeZ conv1_stridemrRrwr2rGsZ           "   zResNet.__init__rT)rVr=r?r>rBcCsjtdd|df|d|}tj|d||df|jdd|j|jd}tj|j|gdd}|S) Navgr!)rCr?r>rr#)dtypedevice)dim) r rzerossizeshaperrcatdata)rOrVr=r?r>rZZ zero_padsr1r1r2_downsample_basic_blocks8zResNet._downsample_basic_block$type[ResNetBlock | ResNetBottleneck]z nn.Sequential)rjr=blocksr>rqr?rBc Csttj|f}ttj|f}d} |dks8|j||jkrt|ddhdkrdt|j||j||d} n.t ||j||jd||j d|||j} ||j|||| dg} ||j|_t d|D]} | ||j||dqt j | S)Nr!r$r")r=r?r>)rCr?rE)r<r=r>r?r@r)rrHr rIr<r`rrrr Sequentialrurangeappend) rOrjr=rr>rqr?rPrQr@rk_ir1r1r2r"sB    zResNet._make_layerrUcCs||}||}||}|js.||}||}||}||}||}| |}| | dd}|j dk r| |}|S)Nr) rJrKrArprrrrrrviewrr)rOrVr1r1r2r[Os           zResNet.forward)r%)r!) r]r^r___doc__rGrrr[rar1r1rRr2rs"! 0G -cs:eZdZdddddddfd d Zd d d dZZS)ResNetFeaturesTr%r!rhrgr7r:) model_name pretrainedr> in_channelsrBc s|tkr,dt}td|d|dt|\}}}} } tj||t||d|d| d |r|dkr|d krt||| d ntd d S) aInitialize resnet18 to resnet200 models as a backbone, the backbone can be used as an encoder for segmentation and objection models. Compared with the class `ResNet`, the only different place is the forward function. Args: model_name: name of model to initialize, can be from [resnet10, ..., resnet200]. pretrained: whether to initialize pretrained MedicalNet weights, only available for spatial_dims=3 and in_channels=1. spatial_dims: number of spatial dimensions of the input image. in_channels: number of input channels for first convolutional layer. z, zinvalid model_name z found, must be one of  r#F) rjrkrlr>rmrorqrtrur%r!) datasets23zQPretrained resnet models are only available for in_channels=1 and spatial_dims=3.N) resnet_paramsjoinkeysrrFrGr3_load_state_dict) rOrrr>rZmodel_name_stringrjrkrqrurrRr1r2rGfs&  zResNetFeatures.__init__rT)inputscCs||}||}||}g}|||js<||}||}||||}||||}||| |}|||S)z Args: inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. Returns: a list of torch Tensors. ) rJrKrArrprrrrr)rOrrVfeaturesr1r1r2r[s             zResNetFeatures.forward)Tr%r!)r]r^r_rGr[rar1r1rRr2rds$rc@sjeZdZdZdddddddgZed d d d Zed d ddZedd ddZedd ddZ dS) ResNetEncoderz9Wrap the original resnet to an encoder for flexible-unet.rrrrrrrz list[dict])rBcCs(g}|jD]}||ddddq |S)z6Get the initialization parameter for resnet backbones.Tr%r!)rrr>r)backbone_namesr)clsparameter_listZ backbone_namer1r1r2get_encoder_parameterss   z$ResNetEncoder.get_encoder_parameterszlist[tuple[int, ...]]cCsdddddddgS)z:Get number of resnet backbone output feature maps channel.)r-r-r.r/r0)r-r/r0iir1rr1r1r2num_channels_per_outputsz%ResNetEncoder.num_channels_per_outputrfcCs dgdS)zGet number of resnet backbone output feature maps. Since every backbone contains the same 5 output feature maps, the number list should be `[5] * 7`. rer1rr1r1r2 num_outputsszResNetEncoder.num_outputsz list[str]cCs|jS)zGet names of resnet backbones.)rrr1r1r2get_encoder_namesszResNetEncoder.get_encoder_namesN) r]r^r_rr classmethodrrrrr1r1r1r2rs  rrhrrfz bool | strrgr)archrjrkrlrprogresskwargsrBcKs|t|||f|}|rxtjr$dnd}t|trjt|r`t d|dtj ||d} nt dn| dddkrP| d dd krF| d d d krFt d|} | rt| d } ntdt| \} } | | ddkr| dkrt| | dd krnt| |d d} n,td| d| dkr4t| ndd| ntdntddd| D} |j| d d|S)NcudacpuzLoading weights from z... map_locationz+The pretrained checkpoint file is not foundr>r%rmr!rtTF resnet(\d+)z1arch argument should be as 'resnet_{resnet_depth}rqr"rrurrzPlease set shortcut_type to z and bias_downsample toz True or Falsez'when using pretrained MedicalNet resnetzgPlease set n_input_channels to 1and feed_forward to False in order to use MedicalNet pretrained weightsz>MedicalNet pretrained weights are only avalaible for 3D modelscSsi|]\}}|dd|qSzmodule.replacervkeyvaluer1r1r2 s z_resnet..)strict)rrr is_availabler~rhrexistsloggerinfoloadFileNotFoundErrorgetresearchr7groupr%get_medicalnet_pretrained_resnet_argsrg get_pretrained_resnet_medicalnetNotImplementedErroritemsload_state_dict)rrjrkrlrrrmodelrmodel_state_dict search_res resnet_depthrurqr1r1r2_resnetsJ    $  &r)rrrrBcKs tdtddddgt||f|S)awResNet-10 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 23 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr!rrr3rrrr1r1r2rs cKs tdtddddgt||f|S)awResNet-18 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 23 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr#rrr1r1r2rs cKs tdtddddgt||f|S)awResNet-34 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 23 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr%r&r'rrr1r1r2r&s cKs tdtddddgt||f|S)awResNet-50 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 23 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr%r&r'rrr3rr1r1r2r2s cKs tdtddddgt||f|S)awResNet-101 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 8 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr%r&r)rrr1r1r2r>s cKs tdtddddgt||f|S)awResNet-152 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 8 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr%r*r+rrr1r1r2rJs cKs tdtddddgt||f|S)awResNet-200 with optional pretrained support when `spatial_dims` is 3. Pretraining from `Med3D: Transfer Learning for 3D Medical Image Analysis `_. Args: pretrained (bool): If True, returns a model pre-trained on 8 medical datasets progress (bool): If True, displays a progress bar of the download to stderr rr%r,r+rrr1r1r2rVs rr7)rrrc Cs d}d}ddddddd g}td ||||kr|sH||d n ||d }zt|||d }Wnxtk r|rt|d|||d }td|t|||d }nt|d||dYnXtj|t|d}ntdt|d| dS)a Download resnet pretrained weights from https://huggingface.co/TencentMedicalNet Args: resnet_depth: depth of the pretrained model. Supported values are 10, 18, 34, 50, 101, 152 and 200 device: device on which the returned state dict will be loaded. "cpu" or "cuda" for example. datasets23: if True, get the weights trained on more datasets (23). Not all depths are available. If not, standard weights are returned. Returns: Pretrained state dict Raises: huggingface_hub.utils._errors.EntryNotFoundError: if pretrained weights are not found on huggingface hub NotImplementedError: if `resnet_depth` is not supported rr "2ez@Loading MedicalNet pretrained model from https://huggingface.co/z.pthz_23dataset.pth)Zrepo_idfilenamez not available for resnetz Trying with z not found on Nrz;Supported resnet_depth are: [10, 18, 34, 50, 101, 152, 200]z downloaded state_dict) rrr Exceptionrrrrrr) rrrZ$medicalnet_huggingface_repo_basenameZ%medicalnet_huggingface_files_basenameZsupported_depthrZpretrained_path checkpointr1r1r2rbsD     r)rcCs(|dkr dnd}|dkrdnd}||fS)z{ Return correct shortcut_type and bias_downsample for pretrained MedicalNet weights according to resnet depth. )rrrrr$r"r1)rrurqr1r1r2rsrz nn.Moduler:)rrrrBcCs`td|}|r*t|d}|d}ntdt|d|d}dd|D}||dS) Nrr!Z _23datasetszZmodel_name argument should contain resnet depth. Example: resnet18 or resnet18_23datasets.rrcSsi|]\}}|dd|qSrrrr1r1r2rs z$_load_state_dict..) rrr7rendswithrrrr)rrrrrrr1r1r2rs  r)FT)FT)FT)FT)FT)FT)FT)rT)T): __future__rloggingrcollections.abcr functoolsrpathlibrtypingrrtorch.nnrZmonai.networks.blocks.encoderrmonai.networks.layers.factoriesrr r Zmonai.networks.layers.utilsr r monai.utilsr monai.utils.modulerrr_rZ$MEDICALNET_HUGGINGFACE_REPO_BASENAMEZ%MEDICALNET_HUGGINGFACE_FILES_BASENAME__all__r getLoggerr]rr3r4Modulerrrrrrrrrrrrrrrrr1r1r1r2 sp         5<-H*8       7