o ) i@@sdZddlmZddlmZmZddlmZddlm Z ddl m Z m Z m Z e d\ZZd gZGd d d ejZ ddddZd 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.modelsBackboneWithFPNcs2eZdZdZ  ddfdd ZdddZZS)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 Nbackbone nn.Module return_layersdict[str, str]in_channels_list list[int] out_channelsint spatial_dims int | None extra_blocksExtraFPNBlock | NonereturnNonecst|dur0t|drt|jtr|j}nt|jtjr"d}nt|jtj r,d}nt d|dur8t |}t j j||d|_t||||d|_||_dS)Nrz;Could not find spatial_dims of backbone, please specify it.)r)rrrr)super__init__hasattr isinstancerrconv1rConv2dConv3d ValueErrorr torchvision_models_utilsZIntermediateLayerGetterbodyr fpnr)selfr rrrrr __class__j/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/backbone_fpn_utils.pyrYs&  zBackboneWithFPN.__init__xrdict[str, Tensor]cCs||}||}|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)r r rrrrrrrrrrrr)r-rrr.)__name__ __module__ __qualname____doc__rr0 __classcell__r+r+r)r,r Cs "Nr resnet.ResNetrrtrainable_layersreturned_layerslist[int] | Nonerrrc s|dks|dkrtd|gdd|}|dkr |d|D]\}tfdd|Dr8|d q$|durAt|}|durIgd }t|dksUt|dkr\td |d d t|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. rr6z3Trainable layers should be in the range [0,5], got )layer4layer3layer2layer1r Nbn1c3s|] }| VqdS)N) startswith).0layer)namer+r, sz(_resnet_fpn_extractor..F)rrrz6Each returned layer should be in the range [1,4]. Got cSs i|] \}}d|t|qS)rB)str)rAvkr+r+r, s z)_resnet_fpn_extractor..csg|] }d|dqS)rrr+)rAi)in_channels_stage2r+r, sz)_resnet_fpn_extractor..)rr) r#appendnamed_parametersallrequires_grad_r minmax enumerateZ in_planesr ) r rr8r9rZlayers_to_train parameterrrrr+)rLrCr,_resnet_fpn_extractors,    rW)r6NN) r r7rrr8rr9r:rrrr )r4 __future__rtorchrrmonai.networks.netsr monai.utilsrfeature_pyramid_networkrr r r$___all__Moduler rWr+r+r+r,s-    J