o ) i)@sdZddlmZddlmZddlmZddlmZddl Z ddl m m Z ddl mZm Z ddlmZmZgd ZGd d d e jZGd d d eZGdddeZGddde jZdS)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/ops/feature_pyramid_network.py ) annotations) OrderedDict)Callable)castN)Tensornn)ConvPool) ExtraFPNBlockLastLevelMaxPool LastLevelP6P7FeaturePyramidNetworkc@seZdZdZd ddZd S) r z Base class for the extra block in the FPN. Same code as https://github.com/pytorch/vision/blob/release/0.12/torchvision/ops/feature_pyramid_network.py results list[Tensor]xnames list[str]cCsdS)av Compute extended set of results of the FPN and their names. Args: results: the result of the FPN x: the original feature maps names: the names for each one of the original feature maps Returns: - the extended set of results of the FPN - the extended set of names for the results Nselfrrrrro/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/feature_pyramid_network.pyforwardKs zExtraFPNBlock.forwardN)rrrrrr)__name__ __module__ __qualname____doc__rrrrrr Dsr cs,eZdZdZdfdd Zdd dZZS)r z Applies a max_pool2d or max_pool3d on top of the last feature map. Serves as an ``extra_blocks`` in :class:`~monai.networks.blocks.feature_pyramid_network.FeaturePyramidNetwork` . spatial_dimsintcs,tttj|f}|dddd|_dS)Nr kernel_sizestridepadding)super__init__r MAXmaxpool)rr pool_type __class__rrr%as zLastLevelMaxPool.__init__rrrrrreturntuple[list[Tensor], list[str]]cCs&|d|||d||fS)Npool)appendr'rrrrrfs zLastLevelMaxPool.forward)rrrrrrrrr+r,rrrrr%r __classcell__rrr)rr [sr cs,eZdZdZdfdd ZdddZZS)r z This module is used in RetinaNet to generate extra layers, P6 and P7. Serves as an ``extra_blocks`` in :class:`~monai.networks.blocks.feature_pyramid_network.FeaturePyramidNetwork` . rr in_channels out_channelscstttj|f}|||dddd|_|||dddd|_|j|jfD]}tjj|j ddtj |j dq&||k|_ dS)Nrrr ar) r$r%rCONVp6p7rinitkaiming_uniform_weight constant_biasuse_P5)rrr3r4 conv_typemoduler)rrr%ss zLastLevelP6P7.__init__rrrrrr+r,c Cs^|d|d}}|jr|n|}||}|t|}|||g|ddg||fS)Nr.r9r:)r@r9r:Freluextend) rrrrp5c5Zx5r9r:rrrr}s zLastLevelP6P7.forward)rrr3rr4rr0r1rrr)rr ls r csDeZdZdZ ddfd d ZdddZdddZdddZZS)r a Module that adds a FPN from on top of a set of feature maps. This is based on `"Feature Pyramid Network for Object Detection" `_. The feature maps are currently supposed to be in increasing depth order. The input to the model is expected to be an OrderedDict[Tensor], containing the feature maps on top of which the FPN will be added. Args: spatial_dims: 2D or 3D images in_channels_list: number of channels for each feature map that is passed to the module out_channels: number of channels of the FPN representation extra_blocks: if provided, extra operations will be performed. It is expected to take the fpn features, the original features and the names of the original features as input, and returns a new list of feature maps and their corresponding names Examples:: >>> m = FeaturePyramidNetwork(2, [10, 20, 30], 5) >>> # get some dummy data >>> x = OrderedDict() >>> x['feat0'] = torch.rand(1, 10, 64, 64) >>> x['feat2'] = torch.rand(1, 20, 16, 16) >>> x['feat3'] = torch.rand(1, 30, 8, 8) >>> # compute the FPN on top of x >>> output = m(x) >>> print([(k, v.shape) for k, v in output.items()]) >>> # returns >>> [('feat0', torch.Size([1, 5, 64, 64])), >>> ('feat2', torch.Size([1, 5, 16, 16])), >>> ('feat3', torch.Size([1, 5, 8, 8]))] Nrrin_channels_list list[int]r4 extra_blocksExtraFPNBlock | Nonec stttj|f}t|_t|_|D]$}|dkr"td|||d}|||ddd}|j ||j |qttj|f} | D] } t | | rhtj j ttj| jddtj ttj| jdqH|durtt |tstt||_dS)Nrz(in_channels=0 is currently not supportedrr5)r#r6g)r$r%rr8r ModuleList inner_blocks layer_blocks ValueErrorr/modules isinstancer;r<rtorchrr=r>r?r AssertionErrorrJ) rrrHr4rJrAr3Zinner_block_moduleZlayer_block_moduleZ conv_type_mr)rrr%s*         zFeaturePyramidNetwork.__init__rridxr+cCFt|j}|dkr ||7}|}t|jD] \}}||kr ||}q|S)zs This is equivalent to self.inner_blocks[idx](x), but torchscript doesn't support this yet r)lenrM enumeraterrrU num_blocksoutirBrrrget_result_from_inner_blocks z2FeaturePyramidNetwork.get_result_from_inner_blockscCrV)zs This is equivalent to self.layer_blocks[idx](x), but torchscript doesn't support this yet r)rWrNrXrYrrrget_result_from_layer_blocksr^z2FeaturePyramidNetwork.get_result_from_layer_blocksdict[str, Tensor]c Cst|}t|}||dd}g}|||dtt|dddD]'}||||}|jdd}t j ||dd} || }| d|||q)|j dur_| |||\}}t tt||} | S)z Computes the FPN for a set of feature maps. Args: x: feature maps for each feature level. Returns: feature maps after FPN layers. They are ordered from highest resolution first. r.rNnearest)sizemoder)listkeysvaluesr]r/r_rangerWshaperC interpolateinsertrJrzip) rrrx_valuesZ last_innerrrUZ inner_lateralZ feat_shapeZinner_top_downr[rrrrs  zFeaturePyramidNetwork.forward)N)rrrHrIr4rrJrK)rrrUrr+r)rr`r+r`) rrrrr%r]r_rr2rrr)rr s+ ! r )r __future__r collectionsrcollections.abcrtypingrrRtorch.nn.functionalr functionalrCrmonai.networks.layers.factoriesrr __all__Moduler r r r rrrrs-