U Ph@@sdZddlmZddlmZmZddlmZddlm Z ddl m Z m Z m Z e d\ZZd gZGd d d ejZdd ddddd dddZd S)z This script is modified from from torchvision to support N-D images, by overriding the definition of convolutional layers and pooling layers. https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py ) annotations)Tensornn)resnet)optional_import) ExtraFPNBlockFeaturePyramidNetworkLastLevelMaxPoolztorchvision.modelsBackboneWithFPNc sDeZdZdZdddddddd d fd d Zd ddddZZS)r a Adds an FPN on top of a model. Internally, it uses torchvision.models._utils.IntermediateLayerGetter to extract a submodel that returns the feature maps specified in return_layers. The same limitations of IntermediateLayerGetter apply here. Same code as https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py Except that this class uses spatial_dims Args: backbone: backbone network return_layers: a dict containing the names of the modules for which the activations will be returned as the key of the dict, and the value of the dict is the name of the returned activation (which the user can specify). in_channels_list: number of channels for each feature map that is returned, in the order they are present in the OrderedDict out_channels: number of channels in the FPN. spatial_dims: 2D or 3D images Nz nn.Modulezdict[str, str]z list[int]intz int | NoneExtraFPNBlock | NoneNone)backbone return_layersin_channels_list out_channels spatial_dims extra_blocksreturncst|dkr`t|dr0t|jtr0|j}n0t|jtjrDd}nt|jtj rXd}nt d|dkrpt |}t j j||d|_t||||d|_||_dS)Nrz;Could not find spatial_dims of backbone, please specify it.)r)rrrr)super__init__hasattr isinstancerr conv1rConv2dConv3d ValueErrorr torchvision_models_utilsZIntermediateLayerGetterbodyr fpnr)selfrrrrrr __class__]/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/backbone_fpn_utils.pyrYs& zBackboneWithFPN.__init__rzdict[str, Tensor])xrcCs||}||}|S)z Computes the resulted feature maps of the network. Args: x: input images Returns: feature maps after FPN layers. They are ordered from highest resolution first. )r"r#)r$r)yr'r'r(forward{s  zBackboneWithFPN.forward)NN)__name__ __module__ __qualname____doc__rr+ __classcell__r'r'r%r(r Cs  "Nz resnet.ResNetr zlist[int] | Noner )rrtrainable_layersreturned_layersrrc s|dks|dkrtd|dddddgd |}|dkrF|d |D](\}tfd d |DrN|d qN|d krt|}|d krddddg}t|dkst|dkrtd|ddt|D}|j dfdd|D}d} t |||| ||dS)a) Same code as https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/backbone_utils.py Except that ``in_channels_stage2 = backbone.in_planes // 8`` instead of ``in_channels_stage2 = backbone.inplanes // 8``, and it requires spatial_dims: 2D or 3D images. rr1z3Trainable layers should be in the range [0,5], got layer4layer3layer2layer1rNbn1c3s|]}| VqdS)N) startswith).0layer)namer'r( sz(_resnet_fpn_extractor..Frrrz6Each returned layer should be in the range [1,4]. Got cSs i|]\}}d|t|qS)r;)str)r:vkr'r'r( sz)_resnet_fpn_extractor..csg|]}d|dqS)rrr')r:i)in_channels_stage2r'r( sz)_resnet_fpn_extractor..)rr) rappendnamed_parametersallrequires_grad_r minmax enumerateZ in_planesr ) rrr2r3rZlayers_to_train parameterrrrr')rEr<r(_resnet_fpn_extractors4    rP)r1NN)r/ __future__rtorchrrmonai.networks.netsr monai.utilsrfeature_pyramid_networkrr r r ___all__Moduler rPr'r'r'r(.s    J