
    Z jk                     <   S r SSKrSSKJr  SSKJr  SSKrSSKJr  SSKJ	r	  SSK
Jr  SS	KJrJr  SS
KJrJr  SSKJr  SSKJrJrJrJr  SSKJrJr  SSKJr  SSKJr  \R@                  " \!5      r"\" SS9\ " S S\5      5       5       r# " S S\RH                  5      r% " S S\RH                  5      r& " S S\RH                  5      r' " S S\RH                  5      r(  S@S\RH                  S\RR                  S\RR                  S \RR                  S!\RR                  S-  S"\*S-  S#\*S$\\   4S% jjr+ " S& S'\RH                  5      r, " S( S)\RH                  5      r- " S* S+\RH                  5      r. " S, S-\RH                  5      r/ " S. S/\RH                  5      r0 " S0 S1\5      r1 " S2 S3\RH                  5      r2\ " S4 S5\5      5       r3\ " S6 S7\35      5       r4 " S8 S9\RH                  5      r5 " S: S;\RH                  5      r6\" S<S9 " S= S>\35      5       r7/ S?Qr8g)AzPyTorch YOLOS model.    N)Callable)	dataclass)nn   )ACT2FN)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPooling)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)ModelOutputTransformersKwargsauto_docstringlogging)can_return_tuplemerge_with_config_defaults)capture_outputs   )YolosConfigz5
    Output type of [`YolosForObjectDetection`].
    )custom_introc                   D   \ rS rSr% SrSr\R                  S-  \S'   Sr	\
S-  \S'   Sr\R                  S-  \S'   Sr\R                  S-  \S'   Sr\\
   S-  \S'   Sr\R                  S-  \S	'   Sr\\R                     S-  \S
'   Sr\\R                     S-  \S'   Srg)YolosObjectDetectionOutput%   a  
loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
    Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
    bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
    scale-invariant IoU loss.
loss_dict (`Dict`, *optional*):
    A dictionary containing the individual losses. Useful for logging.
logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
    Classification logits (including no-object) for all queries.
pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
    Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
    values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
    possible padding). You can use [`~YolosImageProcessor.post_process`] to retrieve the unnormalized bounding
    boxes.
auxiliary_outputs (`list[Dict]`, *optional*):
    Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
    and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
    `pred_boxes`) for each decoder layer.
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
    Sequence of hidden-states at the output of the last layer of the decoder of the model.
Nloss	loss_dictlogits
pred_boxesauxiliary_outputslast_hidden_statehidden_states
attentions )__name__
__module____qualname____firstlineno____doc__r   torchFloatTensor__annotations__r   dictr   r   r   listr    r!   tupler"   __static_attributes__r#       y/root/GenerationalWealth/GenerationalWealth/venv/lib/python3.13/site-packages/transformers/models/yolos/modeling_yolos.pyr   r   %   s    , &*D%

d
")!Itd{!'+FE$++/J!!D(/+/tDzD(/26u((4/659M5**+d2926Je''(4/6r0   r   c                   r   ^  \ rS rSrSrS\SS4U 4S jjrS\R                  S\R                  4S jr	S	r
U =r$ )
YolosEmbeddingsL   zL
Construct the CLS token, detection tokens, position and patch embeddings.

configreturnNc                 x  > [         TU ]  5         [        R                  " [        R
                  " SSUR                  5      5      U l        [        R                  " [        R
                  " SUR                  UR                  5      5      U l	        [        U5      U l        U R                  R                  n[        R                  " [        R
                  " SX!R                  -   S-   UR                  5      5      U l        [        R                  " UR                  5      U l        [#        U5      U l        Xl        g Nr   )super__init__r   	Parameterr)   zeroshidden_size	cls_tokennum_detection_tokensdetection_tokensYolosPatchEmbeddingspatch_embeddingsnum_patchesposition_embeddingsDropouthidden_dropout_probdropout$InterpolateInitialPositionEmbeddingsinterpolationr5   )selfr5   rC   	__class__s      r1   r:   YolosEmbeddings.__init__R   s    ekk!Q8J8J&KL "U[[F<W<WY_YkYk-l m 4V <++77#%<<KK;)D)DDqH&J\J\]$
  zz&"<"<=A&Ir0   pixel_valuesc                 r   UR                   u  p#pEU R                  U5      nUR                  5       u  p'nU R                  R	                  USS5      n	U R
                  R	                  USS5      n
[        R                  " XU
4SS9nU R                  U R                  XE45      nXk-   nU R                  U5      nU$ )Nr   dim)shaperB   sizer>   expandr@   r)   catrI   rD   rG   )rJ   rM   
batch_sizenum_channelsheightwidth
embeddingsseq_len_
cls_tokensr@   rD   s               r1   forwardYolosEmbeddings.forwarda   s    2>2D2D/
&**<8
!+!2
Q ^^**:r2>
0077
BKYY
8HIqQ
 #001I1IF?[5
\\*-
r0   )r>   r5   r@   rG   rI   rB   rD   r$   r%   r&   r'   r(   r   r:   r)   Tensorr^   r/   __classcell__rK   s   @r1   r3   r3   L   s;    
{ t ELL U\\  r0   r3   c                   R   ^  \ rS rSrSU 4S jjrSS\R                  4S jjrSrU =r	$ )rH   v   r6   c                 .   > [         TU ]  5         Xl        g Nr9   r:   r5   rJ   r5   rK   s     r1   r:   -InterpolateInitialPositionEmbeddings.__init__w       r0   c                    US S 2SS S 24   nUS S 2S 4   nUS S 2U R                   R                  * S 2S S 24   nUS S 2SU R                   R                  * 2S S 24   nUR                  SS5      nUR                  u  pgnU R                   R                  S   U R                   R
                  -  U R                   R                  S   U R                   R
                  -  pUR                  XgX5      nUu  pXR                   R
                  -  XR                   R
                  -  p[        R                  R                  X]U4SSS9nUR                  S5      R                  SS5      n[        R                  " X5U4SS9nU$ )Nr   r      bicubicFrS   modealign_cornersrP   )r5   r?   	transposerR   
image_size
patch_sizeviewr   
functionalinterpolateflattenr)   rU   )rJ   	pos_embedimg_sizecls_pos_embeddet_pos_embedpatch_pos_embedrV   r=   r[   patch_heightpatch_widthrX   rY   new_patch_heightnew_patch_widthscale_pos_embeds                   r1   r^   ,InterpolateInitialPositionEmbeddings.forward{   sn   !!Q'*%ag.!!dkk&F&F%F%H!"KL#AqDKK,L,L+L'La$OP)33Aq9+:+@+@(
 KK""1%)?)??KK""1%)?)?? " *..zb ,2kk6L6L,LeWbWbWmWmNm/--33_"EIej 4 
 *11!4>>q!D))]]$SYZ[r0   r5   r6   N)i   i@  
r$   r%   r&   r'   r:   r)   ra   r^   r/   rb   rc   s   @r1   rH   rH   v   s    %,,  r0   rH   c                   R   ^  \ rS rSrSU 4S jjrSS\R                  4S jjrSrU =r	$ ) InterpolateMidPositionEmbeddings   r6   c                 .   > [         TU ]  5         Xl        g rg   rh   ri   s     r1   r:   )InterpolateMidPositionEmbeddings.__init__   rk   r0   c                 R   US S 2S S 2SS S 24   nUS S 2S 4   nUS S 2S S 2U R                   R                  * S 2S S 24   nUS S 2S S 2SU R                   R                  * 2S S 24   nUR                  SS5      nUR                  u  pgpU R                   R                  S   U R                   R
                  -  U R                   R                  S   U R                   R
                  -  pUR                  Xg-  XU5      nUu  pXR                   R
                  -  XR                   R
                  -  p[        R                  R                  X^U4SSS9nUR                  S5      R                  SS5      R                  5       R                  XgX-  U5      n[        R                  " X5U4SS9nU$ )	Nr   r   rm   r   rn   Fro   rP   )r5   r?   rr   rR   rs   rt   ru   r   rv   rw   rx   
contiguousr)   rU   )rJ   ry   rz   r{   r|   r}   depthrV   r=   r[   r~   r   rX   rY   r   r   r   s                    r1   r^   (InterpolateMidPositionEmbeddings.forward   s   !!Q1*-%ag.!!Q)I)I(I(KQ"NO#Aq!t{{/O/O.O*OQR$RS)33Aq92A2G2G/; KK""1%)?)??KK""1%)?)?? " *..u/A;^ij ,2kk6L6L,LeWbWbWmWmNm/--33_"EIej 4 
 ##A&Yq!_Z\T%%5%GU	 	  ))]]$SYZ[r0   r   r   r   r   rc   s   @r1   r   r      s    %,,  r0   r   c                   f   ^  \ rS rSrSrU 4S jrS\R                  S\R                  4S jrSr	U =r
$ )rA      z
This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
`hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
Transformer.
c                   > [         TU ]  5         UR                  UR                  p2UR                  UR
                  pT[        U[        R                  R                  5      (       a  UOX"4n[        U[        R                  R                  5      (       a  UOX34nUS   US   -  US   US   -  -  nX l        X0l        X@l        X`l
        [        R                  " XEX3S9U l        g )Nr   r   )kernel_sizestride)r9   r:   rs   rt   rW   r=   
isinstancecollectionsabcIterablerC   r   Conv2d
projection)rJ   r5   rs   rt   rW   r=   rC   rK   s          r1   r:   YolosPatchEmbeddings.__init__   s    !'!2!2F4E4EJ$*$7$79K9Kk#-j+//:R:R#S#SZZdYq
#-j+//:R:R#S#SZZdYq
!!}
15*Q-:VW=:XY$$(&))L:ir0   rM   r6   c                     UR                   u  p#pEX0R                  :w  a  [        S5      eU R                  U5      R	                  S5      R                  SS5      nU$ )NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.rm   r   )rR   rW   
ValueErrorr   rx   rr   )rJ   rM   rV   rW   rX   rY   rZ   s          r1   r^   YolosPatchEmbeddings.forward   s\    2>2D2D/
&,,,w  __\2::1=GG1M
r0   )rs   rW   rC   rt   r   )r$   r%   r&   r'   r(   r:   r)   ra   r^   r/   rb   rc   s   @r1   rA   rA      s.    jELL U\\  r0   rA   modulequerykeyvalueattention_maskscalingrG   kwargsc                    Uc  UR                  S5      S-  n[        R                  " XR                  SS5      5      U-  nUb  X-   n[        R
                  R                  USS9n[        R
                  R                  XU R                  S9n[        R                  " X5      n	U	R                  SS5      R                  5       n	X4$ )NrO         rm   r   rP   )ptrainingr   )
rS   r)   matmulrr   r   rv   softmaxrG   r   r   )
r   r   r   r   r   r   rG   r   attn_weightsattn_outputs
             r1   eager_attention_forwardr      s     **R.D( <<}}Q':;gEL!#4==((2(>L==((6??([L,,|3K''1-88:K$$r0   c                      ^  \ rS rSrS\4U 4S jjrS\R                  S\\	   S\
\R                  \R                  4   4S jrSrU =r$ )	YolosSelfAttention   r5   c                 0  > [         TU ]  5         UR                  UR                  -  S:w  a7  [	        US5      (       d&  [        SUR                   SUR                   S35      eXl        UR                  U l        [        UR                  UR                  -  5      U l        U R                  U R                  -  U l	        UR                  U l        U R                  S-  U l        SU l        [        R                  " UR                  U R                  UR                   S9U l        [        R                  " UR                  U R                  UR                   S9U l        [        R                  " UR                  U R                  UR                   S9U l        g )	Nr   embedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .r   F)bias)r9   r:   r=   num_attention_headshasattrr   r5   intattention_head_sizeall_head_sizeattention_probs_dropout_probdropout_probr   	is_causalr   Linearqkv_biasr   r   r   ri   s     r1   r:   YolosSelfAttention.__init__   sG    : ::a?PVXhHiHi"6#5#5"6 7334A7 
 #)#=#= #&v'9'9F<V<V'V#W !558P8PP"??//5YYv1143E3EFOO\
99V//1C1C&//ZYYv1143E3EFOO\
r0   r!   r   r6   c                    UR                   S   nUSU R                  U R                  4nU R                  U5      R                  " U6 R                  SS5      nU R                  U5      R                  " U6 R                  SS5      nU R                  U5      R                  " U6 R                  SS5      n[        R                  " U R                  R                  [        5      nU" U UUUS 4U R                  U R                  U R                  (       d  SOU R                   S.UD6u  pU	R#                  5       S S U R$                  4-   nU	R'                  U5      n	X4$ )Nr   rO   r   rm           )r   r   rG   )rR   r   r   r   ru   rr   r   r   r   get_interfacer5   _attn_implementationr   r   r   r   r   rS   r   reshape)rJ   r!   r   rV   	new_shape	key_layervalue_layerquery_layerattention_interfacecontext_layerattention_probsnew_context_layer_shapes               r1   r^   YolosSelfAttention.forward  sO   
 #((+
D$<$<d>V>VV	HH]+00)<FFq!L	jj/44i@JJ1aPjj/44i@JJ1aP(?(M(MKK,,.E)
 *=
*
 nnLL#}}C$2C2C
*
 
*
& #0"4"4"6s";t?Q?Q>S"S%--.EF--r0   )
r   r   r5   r   r   r   r   r   r   r   )r$   r%   r&   r'   r   r:   r)   ra   r   r   r.   r^   r/   rb   rc   s   @r1   r   r      sS    ]{ ](.||. +,. 
u||U\\)	*	. .r0   r   c                      ^  \ rS rSrSrS\4U 4S jjrS\R                  S\R                  S\R                  4S jr	S	r
U =r$ )
YolosSelfOutputi+  z
The residual connection is defined in YolosLayer instead of here (as is the case with other models), due to the
layernorm applied before each block.
r5   c                    > [         TU ]  5         [        R                  " UR                  UR                  5      U l        [        R                  " UR                  5      U l        g rg   )	r9   r:   r   r   r=   denserE   rF   rG   ri   s     r1   r:   YolosSelfOutput.__init__1  sB    YYv1163E3EF
zz&"<"<=r0   r!   input_tensorr6   c                 J    U R                  U5      nU R                  U5      nU$ rg   r   rG   rJ   r!   r   s      r1   r^   YolosSelfOutput.forward6  s$    

=1]3r0   r   r`   rc   s   @r1   r   r   +  sB    
>{ >
U\\  RWR^R^  r0   r   c                   t   ^  \ rS rSrS\4U 4S jjrS\R                  S\\	   S\R                  4S jr
SrU =r$ )	YolosAttentioni=  r5   c                 b   > [         TU ]  5         [        U5      U l        [	        U5      U l        g rg   )r9   r:   r   	attentionr   outputri   s     r1   r:   YolosAttention.__init__>  s&    +F3%f-r0   r!   r   r6   c                 R    U R                   " U40 UD6u  p4U R                  X15      nU$ rg   r   r   )rJ   r!   r   self_attn_outputr\   r   s         r1   r^   YolosAttention.forwardC  s/    
 #nn]EfE-=r0   r   )r$   r%   r&   r'   r   r:   r)   ra   r   r   r^   r/   rb   rc   s   @r1   r   r   =  sC    .{ .
|| +, 
	 r0   r   c                   j   ^  \ rS rSrS\4U 4S jjrS\R                  S\R                  4S jrSr	U =r
$ )YolosIntermediateiN  r5   c                   > [         TU ]  5         [        R                  " UR                  UR
                  5      U l        [        UR                  [        5      (       a  [        UR                     U l        g UR                  U l        g rg   )r9   r:   r   r   r=   intermediate_sizer   r   
hidden_actstrr   intermediate_act_fnri   s     r1   r:   YolosIntermediate.__init__O  s`    YYv1163K3KL
f''--'-f.?.?'@D$'-'8'8D$r0   r!   r6   c                 J    U R                  U5      nU R                  U5      nU$ rg   r   r   )rJ   r!   s     r1   r^   YolosIntermediate.forwardW  s&    

=100?r0   r   r$   r%   r&   r'   r   r:   r)   ra   r^   r/   rb   rc   s   @r1   r   r   N  s/    9{ 9U\\ ell  r0   r   c                      ^  \ rS rSrS\4U 4S jjrS\R                  S\R                  S\R                  4S jrSr	U =r
$ )	YolosOutputi^  r5   c                    > [         TU ]  5         [        R                  " UR                  UR
                  5      U l        [        R                  " UR                  5      U l	        g rg   )
r9   r:   r   r   r   r=   r   rE   rF   rG   ri   s     r1   r:   YolosOutput.__init___  sB    YYv779K9KL
zz&"<"<=r0   r!   r   r6   c                 R    U R                  U5      nU R                  U5      nX-   nU$ rg   r   r   s      r1   r^   YolosOutput.forwardd  s,    

=1]3%4r0   r   r   rc   s   @r1   r   r   ^  s=    >{ >
U\\  RWR^R^  r0   r   c                   x   ^  \ rS rSrSrS\4U 4S jjrS\R                  S\	\
   S\R                  4S jrS	rU =r$ )

YolosLayeril  z?This corresponds to the Block class in the timm implementation.r5   c                 j  > [         TU ]  5         UR                  U l        SU l        [	        U5      U l        [        U5      U l        [        U5      U l	        [        R                  " UR                  UR                  S9U l        [        R                  " UR                  UR                  S9U l        g )Nr   eps)r9   r:   chunk_size_feed_forwardseq_len_dimr   r   r   intermediater   r   r   	LayerNormr=   layer_norm_epslayernorm_beforelayernorm_afterri   s     r1   r:   YolosLayer.__init__o  s    '-'E'E$'/-f5!&) "V-?-?VEZEZ [!||F,>,>FDYDYZr0   r!   r   r6   c                     U R                  U5      nU R                  " U40 UD6nXA-   nU R                  U5      nU R                  U5      nU R	                  XQ5      nU$ rg   )r  r   r  r   r   )rJ   r!   r   hidden_states_normattention_outputlayer_outputs         r1   r^   YolosLayer.forwardy  sl    
 "22=A>>*<GG )8 ++M:((6 {{<?r0   )r   r   r   r  r  r   r   )r$   r%   r&   r'   r(   r   r:   r)   ra   r   r   r^   r/   rb   rc   s   @r1   r   r   l  sH    I[{ [|| +, 
	 r0   r   c                   b   ^  \ rS rSrS\SS4U 4S jjrS\R                  S\S\S\	4S	 jr
S
rU =r$ )YolosEncoderi  r5   r6   Nc                 ^  > [         TU ]  5         Xl        [        R                  " [        UR                  5       Vs/ s H  n[        U5      PM     sn5      U l        SU l	        SUR                  S   UR                  S   -  UR                  S-  -  -   UR                  -   nUR                  (       aD  [        R                  " [        R                   " UR                  S-
  SUUR"                  5      5      OS U l        UR                  (       a  ['        U5      U l        g S U l        g s  snf )NFr   r   rm   )r9   r:   r5   r   
ModuleListrangenum_hidden_layersr   layergradient_checkpointingrs   rt   r?   use_mid_position_embeddingsr;   r)   r<   r=   mid_position_embeddingsr   rI   )rJ   r5   r\   
seq_lengthrK   s       r1   r:   YolosEncoder.__init__  s
   ]]fF^F^@_#`@_1Jv$6@_#`a
&+# ""1%(9(9!(<<@Q@QST@TTUX^XsXss 	 11 LL,,q0&&	  	$ JPIkIk=fEqu' $as   D*r!   rX   rY   c                 V   U R                   R                  (       a  U R                  U R                  X#45      n[	        U R
                  5       HP  u  pVU" U5      nU R                   R                  (       d  M*  XPR                   R                  S-
  :  d  MH  UWU   -   nMR     [        US9$ )Nr   )r    )r5   r  rI   r  	enumerater  r  r	   )rJ   r!   rX   rY   $interpolated_mid_position_embeddingsilayer_modules          r1   r^   YolosEncoder.forward  s     ;;22373E3EdFbFbekds3t0(4OA(7M{{666559:$14XYZ4[$[M  5 ??r0   )r5   r  rI   r  r  )r$   r%   r&   r'   r   r:   r)   ra   r   r	   r^   r/   rb   rc   s   @r1   r
  r
    sS    v{ vt v0@||@ @ 	@
 
@ @r0   r
  c                   N    \ rS rSr% \\S'   SrSrSrSr	/ r
SrSrSrSr\\S.rSrg	)
YolosPreTrainedModeli  r5   vitrM   )imageT)r!   r"   r#   N)r$   r%   r&   r'   r   r+   base_model_prefixmain_input_nameinput_modalitiessupports_gradient_checkpointing_no_split_modules_supports_sdpa_supports_flash_attn_supports_flex_attn_supports_attention_backendr   r   _can_record_outputsr/   r#   r0   r1   r  r    sJ    $O!&*#N"&#(r0   r  c            
          ^  \ rS rSrSS\S\4U 4S jjjrS\4S jr\	\
" SS9\ SS
\R                  S	-  S\\   S\4S jj5       5       5       rSrU =r$ )
YolosModeli  r5   add_pooling_layerc                   > [         TU ]  U5        Xl        [        U5      U l        [        U5      U l        [        R                  " UR                  UR                  S9U l        U(       a  [        U5      OSU l        U R                  5         g)z^
add_pooling_layer (bool, *optional*, defaults to `True`):
    Whether to add a pooling layer
r   N)r9   r:   r5   r3   rZ   r
  encoderr   r   r=   r   	layernormYolosPoolerpooler	post_init)rJ   r5   r+  rK   s      r1   r:   YolosModel.__init__  si    
 	 )&1#F+f&8&8f>S>ST->k&)D 	r0   r6   c                 .    U R                   R                  $ rg   )rZ   rB   )rJ   s    r1   get_input_embeddingsYolosModel.get_input_embeddings  s    ///r0   F)tie_last_hidden_statesNrM   r   c                    Uc  [        S5      eU R                  U5      nUR                  SS  u  pEU R                  X4US9nUR                  nU R                  U5      nU R                  b  U R                  U5      OS n[        XxS9$ )Nz You have to specify pixel_valuesr   )rX   rY   )r    pooler_output)r   rZ   rR   r-  r    r.  r0  r
   )	rJ   rM   r   embedding_outputrX   rY   encoder_outputssequence_outputpooled_outputs	            r1   r^   YolosModel.forward  s     ?@@??<8$**23/+/<<8H_d<+e);;..98<8OO4UY)Oiir0   )r5   rZ   r-  r.  r0  )Trg   )r$   r%   r&   r'   r   boolr:   rA   r4  r   r   r   r)   ra   r   r   r
   r^   r/   rb   rc   s   @r1   r*  r*    s    { t  "0&: 0  E2 -1jllT)j +,j 
$	j  3  jr0   r*  c                   j   ^  \ rS rSrS\4U 4S jjrS\R                  S\R                  4S jrSr	U =r
$ )r/  i  r5   c                    > [         TU ]  5         [        R                  " UR                  UR                  5      U l        [        R                  " 5       U l        g rg   )r9   r:   r   r   r=   r   Tanh
activationri   s     r1   r:   YolosPooler.__init__  s9    YYv1163E3EF
'')r0   r!   r6   c                 \    US S 2S4   nU R                  U5      nU R                  U5      nU$ )Nr   )r   rB  )rJ   r!   first_token_tensorr<  s       r1   r^   YolosPooler.forward  s6     +1a40

#566r0   )rB  r   r   rc   s   @r1   r/  r/    s/    ${ $
U\\ ell  r0   r/  c                   2   ^  \ rS rSrSrU 4S jrS rSrU =r$ )YolosMLPPredictionHeadi  z
Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
height and width of a bounding box w.r.t. an image.

c                    > [         TU ]  5         X@l        U/US-
  -  n[        R                  " S [        U/U-   XS/-   5       5       5      U l        g )Nr   c              3   R   #    U  H  u  p[         R                  " X5      v   M     g 7frg   )r   r   ).0nks      r1   	<genexpr>2YolosMLPPredictionHead.__init__.<locals>.<genexpr>  s     #g@fBIIaOO@fs   %')r9   r:   
num_layersr   r  ziplayers)rJ   	input_dim
hidden_dim
output_dimrP  hrK   s         r1   r:   YolosMLPPredictionHead.__init__  sN    $LJN+mm#gYKRSOUVYeUe@f#ggr0   c                     [        U R                  5       HD  u  p#X R                  S-
  :  a%  [        R                  R                  U" U5      5      OU" U5      nMF     U$ r8   )r  rR  rP  r   rv   relu)rJ   xr  r  s       r1   r^   YolosMLPPredictionHead.forward  sI    !$++.HA01OOa4G0G""58,USTXA /r0   )rR  rP  )	r$   r%   r&   r'   r(   r:   r^   r/   rb   rc   s   @r1   rH  rH    s    h r0   rH  zy
    YOLOS Model (consisting of a ViT encoder) with object detection heads on top, for tasks such as COCO detection.
    c                      ^  \ rS rSrS\4U 4S jjrS r\\ SS\	R                  S\\   S-  S\\   S	\4S
 jj5       5       rSrU =r$ )YolosForObjectDetectioni  r5   c                   > [         TU ]  U5        [        USS9U l        [	        UR
                  UR
                  UR                  S-   SS9U l        [	        UR
                  UR
                  SSS9U l        U R                  5         g )NF)r+  r   r   )rS  rT  rU  rP     )
r9   r:   r*  r  rH  r=   
num_labelsclass_labels_classifierbbox_predictorr1  ri   s     r1   r:    YolosForObjectDetection.__init__!  s      f> (>((V5G5GTZTeTehiTivw(
$ 5((V5G5GTUbc

 	r0   c                 `    [        US S US S 5       VVs/ s H	  u  p4X4S.PM     snn$ s  snnf )NrO   )r   r   )rQ  )rJ   outputs_classoutputs_coordabs        r1   _set_aux_loss%YolosForObjectDetection._set_aux_loss4  s9    ;>}Sb?QS`adbdSe;fg;f411.;fgggs   *NrM   labelsr   r6   c                 ^   U R                   " U40 UD6nUR                  nUSS2U R                  R                  * S2SS24   nU R	                  U5      nU R                  U5      R                  5       nSu  pn
Ub  Su  pU R                  R                  (       a<  UR                  nU R	                  U5      nU R                  U5      R                  5       nU R                  XbU R                  XpR                  X5      u  pn
[        UU	UUU
UR                  UR                  UR                  S9$ )a8  
labels (`list[Dict]` of len `(batch_size,)`, *optional*):
    Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
    following 2 keys: `'class_labels'` and `'boxes'` (the class labels and bounding boxes of an image in the
    batch respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding
    boxes in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image,
    4)`.

Examples:

```python
>>> from transformers import AutoImageProcessor, AutoModelForObjectDetection
>>> import torch
>>> from PIL import Image
>>> import httpx
>>> from io import BytesIO

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> with httpx.stream("GET", url) as response:
...     image = Image.open(BytesIO(response.read()))

>>> image_processor = AutoImageProcessor.from_pretrained("hustvl/yolos-tiny")
>>> model = AutoModelForObjectDetection.from_pretrained("hustvl/yolos-tiny")

>>> inputs = image_processor(images=image, return_tensors="pt")
>>> outputs = model(**inputs)

>>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
>>> target_sizes = torch.tensor([image.size[::-1]])
>>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[
...     0
... ]

>>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
...     box = [round(i, 2) for i in box.tolist()]
...     print(
...         f"Detected {model.config.id2label[label.item()]} with confidence "
...         f"{round(score.item(), 3)} at location {box}"
...     )
Detected remote with confidence 0.991 at location [46.48, 72.78, 178.98, 119.3]
Detected remote with confidence 0.908 at location [336.48, 79.27, 368.23, 192.36]
Detected cat with confidence 0.934 at location [337.18, 18.06, 638.14, 373.09]
Detected cat with confidence 0.979 at location [10.93, 53.74, 313.41, 470.67]
Detected remote with confidence 0.974 at location [41.63, 72.23, 178.09, 119.99]
```N)NNN)NN)r   r   r   r   r   r    r!   r"   )r  r    r5   r?   ra  rb  sigmoidauxiliary_lossr!   loss_functiondevicer   r"   )rJ   rM   rk  r   outputsr;  r   r   r   r   r   re  rf  r   s                 r1   r^   YolosForObjectDetection.forward7  s,   n /3hh|.Nv.N!33 *!dkk.N.N-N-PRS*ST --o>((9AAC
-=**+5(M{{))&44 $ < <\ J $ 3 3L A I I K151C1CZm2.D. *!/%77!//))	
 		
r0   )rb  ra  r  rg   )r$   r%   r&   r'   r   r:   ri  r   r   r)   r*   r-   r,   r   r   r   r^   r/   rb   rc   s   @r1   r]  r]    sw    { &h  %)S
''S
 T
T!S
 +,	S

 
$S
  S
r0   r]  )r]  r*  r  )Nr   )9r(   collections.abcr   r   dataclassesr   r)   r   activationsr   modeling_layersr   modeling_outputsr	   r
   modeling_utilsr   r   processing_utilsr   utilsr   r   r   r   utils.genericr   r   utils.output_capturingr   configuration_yolosr   
get_loggerr$   loggerr   Moduler3   rH   r   rA   ra   floatr   r   r   r   r   r   r   r
  r  r*  r/  rH  r]  __all__r#   r0   r1   <module>r     s7     $ !   ! 9 K F & M M I 5 , 
		H	% 
 7 7 7B'bii 'T299 :ryy B299 P !%II%<<% 
% <<	%
 LL4'% T\% % '(%:4. 4.pbii $RYY "		  
")) 
+ D)@299 )@X ?  " (j% (j (jV"))  RYY & 
l
2 l

l
^ Lr0   