
    \h!                         S SK Jr  S SKJrJr  S SKJs  Jr  S SK	JrJ
r
  SSKJr  SSKJr   " S S	\R                  5      r " S
 S\R                  5      r " S S\5      r " S S\5      rg)    )OrderedDict)CallableOptionalN)nnTensor   )Conv2dNormActivation)_log_api_usage_oncec                   Z    \ rS rSrSrS\\   S\\   S\\   S\\\   \\   4   4S jr	Sr
g	)
ExtraFPNBlock   a~  
Base class for the extra block in the FPN.

Args:
    results (List[Tensor]): the result of the FPN
    x (List[Tensor]): the original feature maps
    names (List[str]): the names for each one of the
        original feature maps

Returns:
    results (List[Tensor]): the extended set of results
        of the FPN
    names (List[str]): the extended set of names for the results
resultsxnamesreturnc                     g )N )selfr   r   r   s       a/var/www/fran/franai/venv/lib/python3.13/site-packages/torchvision/ops/feature_pyramid_network.pyforwardExtraFPNBlock.forward   s     	    r   N__name__
__module____qualname____firstlineno____doc__listr   strtupler   __static_attributes__r   r   r   r   r      sP    f < Cy	
 
tF|T#Y&	'r   r   c                      ^  \ rS rSrSrSr  SS\\   S\S\\	   S\\
S\R                  4      4U 4S	 jjjrU 4S
 jrS\S\S\4S jrS\S\S\4S jrS\\\4   S\\\4   4S jrSrU =r$ )FeaturePyramidNetwork$   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" <https://arxiv.org/abs/1612.03144>`_.

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:
    in_channels_list (list[int]): number of channels for each feature map that
        is passed to the module
    out_channels (int): number of channels of the FPN representation
    extra_blocks (ExtraFPNBlock or None): 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
    norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None

Examples::

    >>> m = torchvision.ops.FeaturePyramidNetwork([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]))]

r   in_channels_listout_channelsextra_blocks
norm_layer.c           
        > [         T	U ]  5         [        U 5        [        R                  " 5       U l        [        R                  " 5       U l        U Hc  nUS:X  a  [        S5      e[        XRSSUS S9n[        X"SUS S9nU R
                  R                  U5        U R                  R                  U5        Me     U R                  5        H  n[        U[        R                  5      (       d  M$  [        R                  R                  UR                  SS9  UR                   c  M[  [        R                  R#                  UR                   S5        M     Ub,  [        U[$        5      (       d  ['        S[)        U5       35      eX0l        g )	Nr   z(in_channels=0 is currently not supported   )kernel_sizepaddingr)   activation_layer   )r,   r)   r.   az1extra_blocks should be of type ExtraFPNBlock not )super__init__r
   r   
ModuleListinner_blockslayer_blocks
ValueErrorr	   appendmodules
isinstanceConv2dinitkaiming_uniform_weightbias	constant_r   	TypeErrortyper(   )
r   r&   r'   r(   r)   in_channelsinner_block_modulelayer_block_modulem	__class__s
            r   r3   FeaturePyramidNetwork.__init__M   sC    	D!MMOMMO+Ka !KLL!5q!PZmq" "6jcg" $$%78$$%78 , A!RYY''((Q(766%GG%%affa0	   #lM::"STXYeTfSg hii(r   c           	      <  > UR                  SS 5      nUb  US:  ak  [        U R                  5      n	S HP  n
[        U	5       H>  nS H5  nU U
 SU SU 3nU U
 SU SU 3nX;   d  M"  UR	                  U5      X'   M7     M@     MR     [
        TU ]  UUUUUUU5        g )Nversionr   )r5   r6   )r>   r?   .z.0.)getlenr5   rangepopr2   _load_from_state_dict)r   
state_dictprefixlocal_metadatastrictmissing_keysunexpected_keys
error_msgsrJ   
num_blocksblockirB   old_keynew_keyrG   s                  r   rP   +FeaturePyramidNetwork._load_from_state_dictp   s     !$$Y5?gkT../J9z*A 2%+HUG1QCq"?%+HUG1QCs4&"A"02<..2IJ/	 !3 + : 	%	
r   r   idxr   c                     [        U R                  5      nUS:  a  X#-  nUn[        U R                  5       H  u  pVXR:X  d  M  U" U5      nM     U$ )z[
This is equivalent to self.inner_blocks[idx](x),
but torchscript doesn't support this yet
r   )rM   r5   	enumerater   r   r^   rX   outrZ   modules          r   get_result_from_inner_blocks2FeaturePyramidNetwork.get_result_from_inner_blocks   T    
 **+
7C"4#4#45IAxQi 6 
r   c                     [        U R                  5      nUS:  a  X#-  nUn[        U R                  5       H  u  pVXR:X  d  M  U" U5      nM     U$ )z[
This is equivalent to self.layer_blocks[idx](x),
but torchscript doesn't support this yet
r   )rM   r6   r`   ra   s          r   get_result_from_layer_blocks2FeaturePyramidNetwork.get_result_from_layer_blocks   rf   r   c                 v   [        UR                  5       5      n[        UR                  5       5      nU R                  US   S5      n/ nUR	                  U R                  US5      5        [        [        U5      S-
  SS5       H`  nU R                  X   U5      nUR                  SS n[        R                  " X7SS9nXh-   nUR                  SU R                  X55      5        Mb     U R                  b  U R                  XAU5      u  pB[        [        X$5       V	V
s/ s H  u  pX4PM
     sn
n	5      nU$ s  sn
n	f )z
Computes the FPN for a set of feature maps.

Args:
    x (OrderedDict[Tensor]): feature maps for each feature level.

Returns:
    results (OrderedDict[Tensor]): feature maps after FPN layers.
        They are ordered from the highest resolution first.
r   Nnearest)sizemoder   )r   keysvaluesrd   r8   rh   rN   rM   shapeFinterpolateinsertr(   r   zip)r   r   r   
last_innerr   r^   inner_lateral
feat_shapeinner_top_downkvrb   s               r   r   FeaturePyramidNetwork.forward   s    QVVX66qubA
t88RHIQ!R,C ==afcJM&,,RS1J]]:YWN&7JNN1d??
PQ - (!..w5ANG c%.AB.AdaA6.ABC
 Cs   D5
)r(   r5   r6   )NN)r   r   r   r   r   _versionr   intr   r   r   r   Moduler3   rP   r   rd   rh   dictr    r   r"   __classcell__rG   s   @r   r$   r$   $   s    $L H 159=!)s)!) !) }-	!)
 Xc299n56!) !)F
@f 3 6 f 3 6  c6k*  tCK/@    r   r$   c                   Z    \ rS rSrSrS\\   S\\   S\\   S\\\   \\   4   4S jr	Sr
g	)
LastLevelMaxPool   z`
Applies a max_pool2d (not actual max_pool2d, we just subsample) on top of the last feature map
r   yr   r   c           	      |    UR                  S5        UR                  [        R                  " US   SSSS95        X4$ )Npoolrk   r+   r   r   )r,   strider-   )r8   rs   
max_pool2d)r   r   r   r   s       r   r   LastLevelMaxPool.forward   s6     	V	ae1aHIxr   r   Nr   r   r   r   r   r      sP    	<	 <	 Cy		
 
tF|T#Y&	'	r   r   c                   z   ^  \ rS rSrSrS\S\4U 4S jjrS\\   S\\   S\\	   S	\
\\   \\	   4   4S
 jrSrU =r$ )LastLevelP6P7   zG
This module is used in RetinaNet to generate extra layers, P6 and P7.
rC   r'   c                   > [         TU ]  5         [        R                  " XSSS5      U l        [        R                  " X"SSS5      U l        U R                  U R
                  4 HU  n[        R                  R                  UR                  SS9  [        R                  R                  UR                  S5        MW     X:H  U l        g )Nr/   r   r+   r0   r   )r2   r3   r   r;   p6p7r<   r=   r>   r@   r?   use_P5)r   rC   r'   rc   rG   s       r   r3   LastLevelP6P7.__init__   s    ))Kq!Q?))L1a@ww(FGG$$V]]a$8GGfkk1- ) "1r   pcr   r   c                     US   US   pTU R                   (       a  UOUnU R                  U5      nU R                  [        R                  " U5      5      nUR                  Xx/5        UR                  SS/5        X4$ )Nrk   r   r   )r   r   r   rs   reluextend)	r   r   r   r   p5c5r   r   r   s	            r   r   LastLevelP6P7.forward   sh     2"B++B2WWQZWWQVVBZ 	"dD\"xr   )r   r   r   )r   r   r   r   r   r   r3   r   r   r    r!   r   r"   r   r   s   @r   r   r      si    2C 2s 2< < Cy	
 
tF|T#Y&	' r   r   )collectionsr   typingr   r   torch.nn.functionalr   
functionalrs   torchr   ops.miscr	   utilsr
   r   r   r$   r   r   r   r   r   <module>r      sU    # %    + 'BII 2hBII hV} "M r   