U Ph:)@sdZddlmZddlmZddlmZddlmm Z ddl m Z mZddl mZmZdd d d gZGd ddejZGd d d eZGdd d eZGdd d ejZdS)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)CallableN)Tensornn)ConvPool ExtraFPNBlockLastLevelMaxPool LastLevelP6P7FeaturePyramidNetworkc@s"eZdZdZddddddZdS)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 list[Tensor] list[str])resultsxnamescCsdS)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 Nselfrrrrrb/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/feature_pyramid_network.pyforwardIs zExtraFPNBlock.forwardN)__name__ __module__ __qualname____doc__rrrrrr Bscs:eZdZdZddfdd Zddddd d d ZZS) 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` . int) spatial_dimscs,tttj|f}|dddd|_dS)Nr kernel_sizestridepadding)super__init__rMAXmaxpool)rr pool_type __class__rrr$_s zLastLevelMaxPool.__init__r rtuple[list[Tensor], list[str]]rrrreturncCs&|d|||d||fS)Npool)appendr&rrrrrds zLastLevelMaxPool.forwardrrrrr$r __classcell__rrr(rr Yscs>eZdZdZddddfdd Zddddd d d ZZS) 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` . r)r in_channels out_channelscstttj|f}|||dddd|_|||dddd|_|j|jfD]&}tjj|j ddtj |j dqL||k|_ dS)Nrrrar) r#r$rCONVp6p7rinitkaiming_uniform_weight constant_biasuse_P5)rrr2r3 conv_typemoduler(rrr$qs zLastLevelP6P7.__init__r rr*r+c Cs^|d|d}}|jr|n|}||}|t|}|||g|ddg||fS)Nr.r8r9)r?r8r9Freluextend) rrrrZp5c5x5r8r9rrrr{s zLastLevelP6P7.forwardr0rrr(rr js csbeZdZdZddddddfdd Zd dd d d d Zd dd d d dZdddddZZS)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]))] Nrz list[int]zExtraFPNBlock | None)rin_channels_listr3 extra_blocksc stttj|f}t|_t|_|D]H}|dkrDtd|||d}|||ddd}|j ||j |q0ttj|f} | D]0} t | | rtj j | jddtj | jdq|dk rt |tst||_dS)Nrz(in_channels=0 is currently not supportedrr4)r"r5g)r#r$rr7r ModuleList inner_blocks layer_blocks ValueErrorr/modules isinstancer:r;r<r=r>r AssertionErrorrH) rrrGr3rHr@r2Zinner_block_moduleZlayer_block_moduleZ conv_type_mr(rrr$s(        zFeaturePyramidNetwork.__init__r)ridxr,cCsFt|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)lenrJ enumeraterrrQZ num_blocksoutirArrrget_result_from_inner_blockss  z2FeaturePyramidNetwork.get_result_from_inner_blockscCsFt|j}|dkr||7}|}t|jD]\}}||kr(||}q(|S)zs This is equivalent to self.layer_blocks[idx](x), but torchscript doesn't support this yet r)rRrKrSrTrrrget_result_from_layer_blockss  z2FeaturePyramidNetwork.get_result_from_layer_blockszdict[str, Tensor])rr,c Cst|}t|}||dd}g}|||dtt|dddD]N}||||}|jdd}t j ||dd} || }| d|||qR|j dk r| |||\}}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)listkeysvaluesrWr/rXrangerRshaperB interpolateinsertrHrzip) rrrx_valuesZ last_innerrrQZ inner_lateralZ feat_shapeZinner_top_downrUrrrrs   zFeaturePyramidNetwork.forward)N) rrrrr$rWrXrr1rrr(rr s +!)r __future__r collectionsrcollections.abcrtorch.nn.functionalr functionalrBtorchrmonai.networks.layers.factoriesrr__all__Moduler r r r rrrr.s