ifddlZddlZddlmZddlmZddlZddlZddl m cm Z ddlm Z m Z ddlmZddlmZddlmZmZmZdd lmZdd lmZmZdd lmZdd lmZm Z m!Z!dd l"m#Z#ddl$m%Z%ddl&m'Z'erddl(m)Z)e!r ddl*m+Z+ddl,m-Z-Gdde j\Z/Gdde j\Z0Gdde j\Z1Gdde j\Z2 dRde j\dejdejdejd ejdzd!e3d"e3fd#Z4Gd$d%e j\Z5dSd&ejd'e3d(e6d)ejfd*Z7Gd+d,e j\Z8Gd-d.e j\Z9Gd/d0eZ:Gd1d2e j\Z;ee d34Gd5d6eZ< dTd7ejd8ejd)ejfd9Z=d:e d;e d)e fd<Z>d:ejd;ejd)ejfd=Z?Gd>d?e j\Z@d:e d;e d@eAd)e fdAZBd:ejd;ejd@eAd)ejfdBZCGdCdDe j\ZDe GdEdFeZEGdGdHe jZGGdIdJe j\ZHGdKdLe j\ZIGdMdNe j\ZJe dO4GdPdQeEZKdFdQgZLy)UN)Callable) dataclass)Tensornn)initialization)ACT2FN) ModelOutputis_scipy_availablerequires_backends)GradientCheckpointingLayer)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)TransformersKwargsauto_docstringis_accelerate_available)merge_with_config_defaults)capture_outputs) VideomtConfig)linear_sum_assignment) PartialState)reducecZeZdZdZfdZdej dej fdZxZS)VideomtPatchEmbeddingsz 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. ct||j|j}}|j|j }}t |tjjr|n||f}t |tjjr|n||f}|d|dz|d|dzz}||_||_||_||_ tj|||||_ y)Nrr kernel_sizestride)super__init__ image_size patch_size num_channels hidden_size isinstance collectionsabcIterable num_patchesrConv2d projection)selfconfigr#r$r%r&r+ __class__s }/var/www/vps2.regionflexible.com/Desarrollo/venv/lib/python3.12/site-packages/transformers/models/videomt/modeling_videomt.pyr"zVideomtPatchEmbeddings.__init__:s !'!2!2F4E4EJ $*$7$79K9Kk #-j+//:R:R#SZZdfpYq #-j+//:R:R#SZZdfpYq !!} 1 5*Q-:VW=:XY $$(&))L+:^hi pixel_valuesreturnc:|jd}||jk7rtd|jd|d|j|jj j }|j |jdjdd}|S)NrzoMake sure that the channel dimension of the pixel values match with the one set in the configuration. Expected z but got .dtype) shaper% ValueErrortor-weightr8flatten transpose)r.r3r% embeddingss r1forwardzVideomtPatchEmbeddings.forwardIs#))!, 4,, ,!../yaI  $T__-C-C-I-IJ __\2::1=GG1M r2) __name__ __module__ __qualname____doc__r"torchrrA __classcell__r0s@r1rr3s) j ELL U\\ r2rceZdZdZdeddffd Zd dejdejdzdejfdZxZ S) VideomtEmbeddingszM Construct the CLS token, mask token, position and patch embeddings. r/r4Nc.t|||_|j|_t j t jdd|j|_ t j t jd|j|j|_ t||_|jj}t j |j"|_d|jz|_t j(||j|_|j-dt j.|j1ddt j t jdd|j|_y)Nr position_idsrF) persistent)r!r"r/r$r ParameterrFrandnr& cls_tokenzerosnum_register_tokensregister_tokensrpatch_embeddingsr+Dropouthidden_dropout_probdropoutnum_prefix_tokens Embeddingposition_embeddingsregister_bufferarangeexpand mask_token)r.r/r+r0s r1r"zVideomtEmbeddings.__init__[s   ++ekk!Q8J8J&KL!||EKK6;U;UW]WiWi,jk 6v >++77 zz&"<"<= !"V%?%?!?#%<< V=O=O#P  ^U\\+-F-M-Mg-Vchi,,u{{1a9K9K'LMr2r3bool_masked_posc|jdk(rA|j\}}}}}|j||z|||}|F|j||zd}n0|.|jdkDr|j|jdd}|jd}|j|}|[|j |j t jjd} t j| |j|}|jj|dd} |jj|dd} ||j|jz}t j | | |gd}|j#|}|S)NrNr9r)devicer8rdim)ndimr:reshaperVr<rdrFbool unsqueezewherer`rRr_rUr\rLcatrY) r.r3ra batch_size num_framesr%heightwidthr@mask cls_tokensrUs r1rAzVideomtEmbeddings.forwardksq    !BNBTBT ?J L&%'// Z0GW]_deL*"1"9"9*z:QSU"V  (_-A-AA-E-55o6K6KA6NPRSO!''* **<8  &"%%Z->->ejj%Q[[\^_DT4??JGJ^^**:r2> ..55j"bI$":":4;L;L"MM YY OZHaP \\*- r2N rBrCrDrErr"rFrrArGrHs@r1rJrJVsMN}NN ELL5<> )T^^ ;MdnnM]^NN#2' ]]D( // ii? ii? ii?  $..$..A r2N hidden_statesrr4c |jdd}g|d|j}|j|j|j dd}|j |j|j dd}|j |j|j dd}tj|jjt} | ||||||j|j|jsdn |j\} } | j g|dj#} |j%| } | | fS)z#Input shape: Batch x Time x ChannelNrNrr9)rrrY)r:rrviewr?rrr get_interfacer/_attn_implementationrrrrrYrhrr) r.rrr input_shape hidden_shapequerieskeysvaluesattention_interfacerrs r1rAzVideomtAttention.forwards<$))#2. 88b8$--8 ++m,11,?II!QO{{=)..|<FFq!L]+00>HHAN(?(M(M KK , ,.E) %8     nnJJ#}}C$,, % ! \*k));;;;FFH mmK0 L((r2rs) rBrCrDrEr"rFrtuplerArGrHs@r1rrsVGB./3!)||!) t+!) u||U\\D00 1 !)r2rinput drop_probrr4c|dk(s|s|Sd|z }|jdfd|jdz zz}|tj||j|j z}|j |j||z}|S)zc Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). rrr)rr8rd)r:rgrFrandr8rdfloor_div)rrr keep_probr: random_tensoroutputs r1 drop_pathrs Cx I I [[^  Q 7 7E 5 ELL YYM YYy !M 1F Mr2cxeZdZdZd dedzddffd ZdejdejfdZde fdZ xZ S) VideomtDropPathzXDrop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).Nrr4c0t|||_yrs)r!r"r)r.rr0s r1r"zVideomtDropPath.__init__ s "r2rcDt||j|jSrs)rrrr.rs r1rAzVideomtDropPath.forwards FFr2c d|jS)Nzp=)rr.s r1 extra_reprzVideomtDropPath.extra_reprsDNN#$$r2rs) rBrCrDrEfloatr"rFrrArrrGrHs@r1rr sGb#%$,#$#GU\\GellG%C%r2rcXeZdZdfd Zdej dej fdZxZS)VideomtSwiGLUFFNr4c0t||jx}}t|j|jz}t|dzdz dzdzdz}t j |d|zd|_t j ||d|_yrrrs r1r"zVideomtSwiGLUFFN.__init__rr2rc|j|}|jdd\}}tjj ||z}|j |Srrrs r1rAzVideomtSwiGLUFFN.forward"rr2rrrHs@r1rrrr2rceZdZdZdeddffd Z d dejdejdzdejfdZxZ S) VideomtLayerzCThis corresponds to the Block class in the original implementation.r/r4Nct|tj|j|j |_t||_t||_ |jdkDrt|jntj|_tj|j|j |_|j rt#||_nt'||_t||_y)Nepsr)r!r"r LayerNormr&layer_norm_epsnorm1r attentionVideomtLayerScale layer_scale1drop_path_raterIdentityrnorm2use_swiglu_ffnrmlprv layer_scale2rs r1r"zVideomtLayer.__init__,s \\&"4"4&:O:OP )&1-f5CICXCX[^C^)>)>?dfdododq\\&"4"4&:O:OP  '/DH!&)DH-f5r2rrc*|j|}|j||\}}|j|}|j||z}|j |}|j |}|j |}|j||z}|Srs)rrrrrrr)r.rrhidden_states_normself_attention_output_ layer_outputs r1rAzVideomtLayer.forward<s "ZZ 6#'>>2Dn#U q $ 1 12G H'<= M zz-0 xx - ((6 ~~l3mC r2rsrtrHs@r1rr)sPM6}66&/3|| t+  r2rcXeZdZdfd Zdej dej fdZxZS)rr4ct|tj|jt j |jz|_yrs) r!r"rrPlayerscale_valuerFonesr&lambda1rs r1r"zVideomtLayerScale.__init__Ts8 ||F$;$;ejjI[I[>\$\] r2rc ||jzSrs)rrs r1rAzVideomtLayerScale.forwardXsdll**r2rrrHs@r1rrSs$^+ELL+U\\+r2ra Class for outputs of [`VideomtForUniversalSegmentationOutput`]. This output can be directly passed to [`~VideomtVideoProcessor.post_process_semantic_segmentation`] or [`~VideomtVideoProcessor.post_process_instance_segmentation`] or [`~VideomtVideoProcessor.post_process_panoptic_segmentation`] to compute final segmentation maps. Please, see [`~VideomtVideoProcessor`] for details regarding usage. ) custom_introceZdZUdZdZej dzed<dZej dzed<dZ ej dzed<dZ ej dzed<dZ e ej dzed<dZ e ej dzed<y) %VideomtForUniversalSegmentationOutputa loss (`torch.Tensor`, *optional*): The computed loss, returned when labels are present. class_queries_logits (`torch.FloatTensor`): A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each query. Note the `+ 1` is needed because we incorporate the null class. masks_queries_logits (`torch.FloatTensor`): A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each query. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Last hidden states (final feature map) of the last layer. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states all layers of the model. attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Self and Cross Attentions weights from transformer decoder. Nlossclass_queries_logitsmasks_queries_logitslast_hidden_stater attentions)rBrCrDrErrF FloatTensor__annotations__rrr rrr r2r1rr\s&&*D%  d ")59%++d2959%++d2926u((4/659M5**+d2926Je''(4/6r2rinput_featurespoint_coordinatesc |jdk(rd}|jd}tjjj |d|zdz fi|}|r|j d}|S)a( A wrapper around `torch.nn.functional.grid_sample` to support 3D point_coordinates tensors. Args: input_features (`torch.Tensor` of shape (batch_size, channels, height, width)): A tensor that contains features map on a height * width grid point_coordinates (`torch.Tensor` of shape (batch_size, num_points, 2) or (batch_size, grid_height, grid_width,: 2)): A tensor that contains [0, 1] * [0, 1] normalized point coordinates add_dim (`bool`): boolean value to keep track of added dimension Returns: point_features (`torch.Tensor` of shape (batch_size, channels, num_points) or (batch_size, channels, height_grid, width_grid): A tensor that contains features for points in `point_coordinates`. rTr9g@?)rfrjrFrr grid_samplesqueeze)rradd_dimrpoint_featuress r1 sample_pointrst(!#-77:XX((44^SK\E\_bEbmflmN'//2 r2inputslabelsc|jjd}dtj||jz}|j ddddf|j ddddfz}d|dz|dzz z }|S)a A pair wise version of the dice loss, see `dice_loss` for usage. Args: inputs (`torch.Tensor`): A tensor representing a mask labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). Returns: `torch.Tensor`: The computed loss between each pairs. rr9rNN)sigmoidr>rFrTsum)rr numerator denominatorrs r1pair_wise_dice_lossrs|^^  % %a (FELL22I**R.D)FJJrN47,CCK A +/2 2D Kr2c\|jd}tjd}||tj|}||tj |}tj ||z |j}tj ||z d|z j}||z}|S)a A pair wise version of the cross entropy loss, see `sigmoid_cross_entropy_loss` for usage. Args: inputs (`torch.Tensor`): A tensor representing a mask. labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). Returns: loss (`torch.Tensor`): The computed loss between each pairs. rnone reduction)r:rBCEWithLogitsLossrF ones_like zeros_likerr) rrheight_and_width criterioncross_entropy_loss_poscross_entropy_loss_negloss_posloss_negrs r1$pair_wise_sigmoid_cross_entropy_lossr-s||A$$v6I&vuv/FG&vu/?/?/GH||25EEvxxPH||25EEF ~~VH h D Kr2c eZdZdZ d dedededeffd Zejdejdejd ejd ejd e e e f d Z xZ S)VideomtHungarianMatcheraqThis class computes an assignment between the labels and the predictions of the network. For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more predictions than labels. In this case, we do a 1-to-1 matching of the best predictions, while the others are un-matched (and thus treated as non-objects). cost_class cost_mask cost_dice num_pointsct||dk(r|dk(r|dk(r td||_||_||_||_y)aHCreates the matcher Params: cost_class (`float`, *optional*, defaults to 1.0): Relative weight of the classification error in the matching cost. cost_mask (`float`, *optional*, defaults to 1.0): This is the relative weight of the focal loss of the binary mask in the matching cost. cost_dice (`float`, *optional*, defaults to 1.0): This is the relative weight of the dice loss of the binary mask in the matching cost. num_points (`int`, *optional*, defaults to 12544): No. of points to sample on which the mask loss will be calculated. The same set of K points are uniformly sampled for all prediction and ground truth masks to construct the cost matrix for bipartite matching. rzAll costs can't be 0N)r!r"r;r3r0r1r2)r.r0r1r2r3r0s r1r"z VideomtHungarianMatcher.__init__sK"  ?yA~)q.34 4$$""r2rr mask_labels class_labelsr4c g}|jd}t|D]}||jd}||} |dd||f } ||j| } | dddf} | dddf} t j d|j d| j} | j| jddd} t| | djd} | j| jddd}t| |djd} t| | }t| | }|j|z|j| zz|j|zz}t j |t j"d }t j$|t j"d }t j&|d}t)|j+}|j-||Dcgc]O\}}t j.|tj0 t j.|tj0 fQ}}}|Scc}}w) ao Params: masks_queries_logits (`torch.Tensor`): A tensor of dim `batch_size, num_queries, num_labels` with the classification logits. class_queries_logits (`torch.Tensor`): A tensor of dim `batch_size, num_queries, height, width` with the predicted masks. class_labels (`torch.Tensor`): A tensor of dim `num_target_boxes` (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels. mask_labels (`torch.Tensor`): A tensor of dim `num_target_boxes, height, width` containing the target masks. Returns: matched_indices (`list[tuple[Tensor]]`): A list of size batch_size, containing tuples of (index_i, index_j) where: - index_i is the indices of the selected predictions (in order) - index_j is the indices of the corresponding selected labels (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes). rrNNrr9rdF align_cornersg _Bg _r7)r:rangerr<rFrr3rdrepeatrrr-rr1r0r2minimumtensormaximum nan_to_numrcpuappend as_tensorint64)r.rrr5r6indicesrmi pred_probs pred_maskr0 target_maskrtarget_coordinatespred_coordinatesr1r2 cost_matrixassigned_indicesjmatched_indicess r1rAzVideomtHungarianMatcher.forwards/8*,*//2 z" -A-a088r)rrrRprobsrrrs r1 dice_lossrU@sz, NN  $ $Q 'EUV^((,,I))B-&**R.0K A +/2 2D 88: !D Kr2ctjd}|||}|jdj|z }|S)a| Args: inputs (`torch.Tensor`): A float tensor of arbitrary shape. labels (`torch.Tensor`): A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs (0 for the negative class and 1 for the positive class). Returns: loss (`torch.Tensor`): The computed loss. r!r"r)rr$meanr)rrrRr(cross_entropy_lossrs r1sigmoid_cross_entropy_lossrY^sD$$v6I"662  " "1 % ) ) +i 7D Kr2ceZdZdedeeefffd Zdeee dee fdZ dee de e e ffdZ d e d ee d e ejdeee ffd Zd ejdeejd e ejde deeejff dZdZdZdejdejfdZdejde de dedejf dZ dd ejd ejdeejd eejdeeejfdzdeeejff dZd ejdej0dejfdZxZS) VideomtLossr/ weight_dictct|t|dg|j|_||_|j |_tj|jdz}|j |d<|jd||j|_ |j|_ |j|_ t|j|j |j"|j|_y)aQ The Videomt Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we compute hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair of matched ground-truth / prediction (supervise class and mask) Args: config (`VideomtConfig`): The configuration for Videomt model also containing loss calculation specific parameters. weight_dict (`dict[str, float]`): A dictionary of weights to be applied to the different losses. scipyrrN empty_weight)r0r2r1r3N)r!r"r num_labelsr\no_object_weighteos_coefrFrr]train_num_pointsr3oversample_ratioimportance_sample_ratior/ class_weight dice_weight mask_weightmatcher)r.r/r\r_r0s r1r"zVideomtLoss.__init__ss $ * ++&// zz$//A"56 == R ^\:!11 & 7 7'-'E'E$.**((((   r2sizesr4cn|d}|ddD]'}t|D]\}}t|||||<)|S)Nrr) enumeratemax)r.rjmaxessublistindexitems r1 _max_by_axiszVideomtLoss._max_by_axissSaQRy 7G(1 7 t"5<6e  7 7 r2tensorsc`|j|Dcgc]}t|jc}}t|g|z}|\}}}}|dj} |dj } t j|| | } t j|||ft j| } t|| | D]o\}} }| d|jdd|jdd|jdfj|d|d|jdd|jdf<q| | fScc}w)Nrrrr9F) rrrQr:lenr8rdrFrSrrizipcopy_)r.rsr>max_size batch_shapermrrorpr8rdpadded_tensors padding_masks padded_tensor padding_masks r1_pad_images_to_max_in_batchz'VideomtLoss._pad_images_to_max_in_batchs7$$w%OVd6<<&8%OP7|nx/ '2$ Avu   ""[fM J#>ejjY_` 36wP]3^ G /FM< +FLLO+->v||A->@Q&,,q/@QQ R X XY_ `AFL*6<<?*,=fll1o,== > G},,&PsD+rr6rEc |}|j\}}}tj|j}|j |} t j t||D  cgc] \} \}} | | c} }} } t j||f|jt j|j} | | | <|jdd}||| }d|i}|Scc} }} w)aCompute the losses related to the labels using cross entropy. Args: class_queries_logits (`torch.Tensor`): A tensor of shape `batch_size, num_queries, num_labels` class_labels (`list[torch.Tensor]`): List of class labels of shape `(labels)`. indices (`tuple[np.array])`: The indices computed by the Hungarian matcher. Returns: `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key: - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. )r=) fill_valuer8rdrr9loss_cross_entropy) r:rCrossEntropyLossr_$_get_predictions_permutation_indicesrFrlrvfullr`rDrdr?)r.rr6rE pred_logitsrm num_queriesrr(idxtargetrNtarget_classes_otarget_classespred_logits_transposedloss_celossess r1 loss_labelszVideomtLoss.loss_labelss"+ %0%6%6" K''t/@/@A 77@ 99-0w-G H H>66AqVAY H   %$//]h]o]o /s!,!6!6q!!<2NC&0  Is#C! rr5rRcj|}j|}||}j|\}} ||}|dddf}|dddf}tj5j |fdj jj} t|| djd} dddt| djd} t|  |t| | |d} ~~| S#1swYExYw)aCompute the losses related to the masks using sigmoid_cross_entropy_loss and dice loss. Args: masks_queries_logits (`torch.Tensor`): A tensor of shape `(batch_size, num_queries, height, width)`. mask_labels (`torch.Tensor`): List of mask labels of shape `(labels, height, width)`. indices (`tuple[np.array])`: The indices computed by the Hungarian matcher. num_masks (`int)`: The number of masks, used for normalization. Returns: losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing two keys: - **loss_mask** -- The loss computed using sigmoid cross entropy loss on the predicted and ground truth. masks. - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth, masks. Nc&j|Srs)calculate_uncertainty)logitsr.s r1z(VideomtLoss.loss_masks..st99&Ar2Fr9r) loss_mask loss_dice) r _get_targets_permutation_indicesr~rFrPsample_points_using_uncertaintyr3rdrerrrYrU)r.rr5rErRsrc_idxtgt_idx pred_masks target_masksrr point_labels point_logitsrs` r1 loss_maskszVideomtLoss.loss_maskss,4;;GD77@)'2 ::;G a#G,  4( #AtG, ]]_ i $ D DA%%,, ! ( 6GW\]eefghL i$J0AQVW__`ab 4L,PYZ"<yI    ) i is (ADD c tjt|Dcgc]\}\}}tj|| c}}}}tj|Dcgc]\}}| c}}}||fScc}}}wcc}}wrsrFrlrl full_like)r.rErFsrcr batch_indicespredictions_indicess r1rz0VideomtLoss._get_predictions_permutation_indicessj iX_N`"a"a{q(35??3#:"ab #iiW(E#q(EF111#b(E #A7  A> c tjt|Dcgc]\}\}}tj|| c}}}}tj|Dcgc]\}}| c}}}||fScc}}}wcc}}wrsr)r.rErFrtgtrtarget_indicess r1rz,VideomtLoss._get_targets_permutation_indicessh iX_N`"a"a{q(1c5??3#:"ab #@HQC#@An,,#b#@rrc2tj| }|S)a In Videomt paper, uncertainty is estimated as L1 distance between 0.0 and the logit prediction in 'logits' for the foreground class in `classes`. Args: logits (`torch.Tensor`): A tensor of shape (R, 1, ...) for class-specific or class-agnostic, where R is the total number of predicted masks in all images and C is: the number of foreground classes. The values are logits. Returns: scores (`torch.Tensor`): A tensor of shape (R, 1, ...) that contains uncertainty scores with the most uncertain locations having the highest uncertainty score. )rFabs)r.runcertainty_scoress r1rz!VideomtLoss.calculate_uncertaintys %yy01!!r2r3rdrec |jd}t||z}tj||d|j}t ||d} || } t||z} || z } tj | dddddf| dd} |tj|tj|j z}| |dddfz } |jd d| jd ddfj|| d}| dkDr:tj|tj|| d|jgd }|S) a This function is meant for sampling points in [0, 1] * [0, 1] coordinate space based on their uncertainty. The uncertainty is calculated for each point using the passed `uncertainty function` that takes points logit prediction as input. Args: logits (`float`): Logit predictions for P points. uncertainty_function: A function that takes logit predictions for P points and returns their uncertainties. num_points (`int`): The number of points P to sample. oversample_ratio (`int`): Oversampling parameter. importance_sample_ratio (`float`): Ratio of points that are sampled via importance sampling. Returns: point_coordinates (`torch.Tensor`): Coordinates for P sampled points. rr9r8Fr9Nr)krfrrNre) r:rzrFrrdrtopkr^longrrl)r.runcertainty_functionr3rdre num_boxesnum_points_sampledrrpoint_uncertaintiesnum_uncertain_pointsnum_random_pointsrshifts r1rz+VideomtLoss.sample_points_using_uncertainty,sI<LLO  .>!>?"JJy2DaPVP]P]^#F,=US 2<@"#:Z#GH&)==jj,Q1W59MSTUVWX"U\\)5::V\VcVc%dd uQW~-222q9#((2,/JOOPY[oqrs q % "EJJy:KQW]WdWd$ef! ! r2Nauxiliary_predictionsc|j||||}|j||dj}i|j|||||j |||}|jt |D]\\} } | d}| d}|j ||||} | jD cic] \} } | d| | } } } |j| ^|Scc} } w)a This performs the loss computation. Args: masks_queries_logits (`torch.Tensor`): A tensor of shape `(batch_size, num_queries, height, width)`. class_queries_logits (`torch.Tensor`): A tensor of shape `(batch_size, num_queries, num_labels)`. mask_labels (`torch.Tensor`): List of mask labels of shape `(labels, height, width)`. class_labels (`list[torch.Tensor]`): List of class labels of shape `(labels)`. auxiliary_predictions (`dict[str, torch.Tensor]`, *optional*): if `use_auxiliary_loss` was set to `true` in [`VideomtConfig`], then it contains the logits from the inner layers of the VideomtMaskedAttentionDecoder. Returns: losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing three keys: - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels. - **loss_mask** -- The loss computed using sigmoid cross_entropy loss on the predicted and ground truth masks. - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth masks. if `use_auxiliary_loss` was set to `true` in [`VideomtConfig`], the dictionary contains additional losses for each auxiliary predictions. rr8rrr) ri get_num_masksrdrrrlrAitemsupdate)r.rrr5r6rrErRrr aux_outputs loss_dictrrs r1rAzVideomtLoss.forwardcs H,,35I;Xde&&|LOz,VideomtLoss.get_num_masks..sAG Asrr)min) rrFrCrrr _shared_stater num_processesclamp)r.r6rdrR world_sizes r1rzVideomtLoss.get_num_maskssuALAA OOIU[[P  " $))R/"9- )^99 KK J 6A> r2rs)rBrCrDrdictrrr"rQrzrrrrr~nparrayrrFrrrrrrArdrrGrHs@r1r[r[rs! }! 4U ;K! F$tCy/d3i-4<-E&RX.DY-" $* :>v, QVWYW_W_Q`  c6k  D<#ll<%,,'<rxx <  < c5<< <|2 - "ELL"U\\""5! 5! 5!  5! "' 5! 5!zAE 5#ll5$ll5%,,' 5 5<<( 5 $C$56= 5 c5<< 5n %,,  QVQ]Q] r2r[ceZdZUdZeed<dZdZdZdZ dgZ dZ e e d Zej d ej$d d fd Zy )VideomtPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. r/videomtpixel_values_videos)videoFrT)rr rr4Ncz|jj}t|tjtj tj frtj|jtjd|jtjjj|j\}}|dkDrdtj|z nd}tj|j| |nxt|tj r@tj"|jtj$|jnt|tj&rutj(|jdd|j*t-|jddstj$|j|j*nt|t.rCt1|d rrtj2|j4|jj6ntj@|jBtjD|jBjFd jId nt|tJrRtjL|jNdz}|jP|d <tj@|jR|n/t|tTrtj"|jVt|t8r*tjj%|jXyy) Nrc)arrr)rWstd_is_hf_initializedFrrNrM)-r/initializer_ranger'rr|r,ConvTranspose2dinitkaiming_uniform_r=mathsqrtryrF_calculate_fan_in_and_fan_outuniform_rones_zeros_r[normal_ padding_idxgetattrrhasattr constant_rrrJ trunc_normal_rRrUrwrLr^r:r_r[rr`rbr_VideomtForUniversalSegmentationattn_mask_probsr`)r.rrfan_inrboundr_s r1 _init_weightsz$VideomtPreTrainedModel._init_weightssXkk++ fryy"))R5G5GH I  ! !&--499Q< @{{&!HHMMGG V 17!DIIf-- fkkE659  - JJv}} % KK $  - LLSa 8!!-gfmmMach6i FMM&*<*<=>  1 2vy)v~~t{{/K/KL  1 2   v//cs C KK.. / JJv**ELL9L9L9R9RSU9V,W,^,^_f,g h  , ::f&7&7!&;TrrHs@r1rrs$KELLU\\r2rc\eZdZdeffd Zdej dej fdZxZS)VideomtScaleLayerr/ct||j}tj||dd|_t |j|_tj||dd|d|_ t||_ y)Nr9rrrF)rpaddinggroupsry) r!r"r&rrconv1r r~rr,conv2r layernorm2dr.r/r&r0s r1r"zVideomtScaleLayer.__init__su (( '' [aXYZ  !2!23YY     .k:r2rr4c|j|}|j|}|j|}|j|}|Srs)rrrrrs r1rAzVideomtScaleLayer.forwardsB =1  6  =1 ((7 r2 rBrCrDrr"rFrrArGrHs@r1rrs*;}; U\\ellr2rc\eZdZdeffd Zdej dej fdZxZS)VideomtScaleBlockr/ct||j|_t j t |jDcgc] }t|c}|_ycc}wrs) r!r"num_upscale_blocks num_blocksr ModuleListr;rblockr.r/rr0s r1r"zVideomtScaleBlock.__init__sH  33]]uT__G]#^!$5f$=#^_ #^sA&rr4c8|jD] }||} |Srs)r )r.rr s r1rAzVideomtScaleBlock.forward s%ZZ 1E!-0M 1r2rrHs@r1rrs,`}` U\\ellr2rc\eZdZdeffd Zdej dej fdZxZS)VideomtMaskHeadr/c t||j}tj|||_tj|||_tj|||_t|j|_ yrs) r!r"r&rr|r}rfc3r r~rrs r1r"zVideomtMaskHead.__init__sa (( 99[+699[+699[+6 !2!23r2rr4c|j|j|}|j|j|}|j|}|Srs)rr}rrrs r1rAzVideomtMaskHead.forwardsD(?@ (?@ / r2rrHs@r1rrs*4}4U\\ellr2rzY The Videomt Model with head on top for instance/semantic/panoptic segmentation. cdeZdZdZdeffd Zdededededeeefd eeeff d Z d eeefd efd Z e e e ddej d zdeej d zdeej d zdeej d zdeed ef dZdZdej fdZxZS)rrr/c<t||||_|j|_t ||_t j|j|j|_ t j|j|j|_ t jt|jDcgc] }t!|c}|_t%||_t)||_t j,|j|j.dz|_|j2|j4z|j2|j4zf|_|j8|j:|j<d|_tA||j>|_!|jEdtGjH|jJt j,|j|j|_&|jOycc}w)Nrr)rrr)r/r\r)(r!r"r/num_hidden_layersrJr@rrr&r layernormr[rrr r;rlayersr upscale_blockr mask_headr|r`class_predictorr#r$ grid_sizerfrhrgr\r[r(r]rFrr  query_updater post_initr s r1r"z(VideomtForUniversalSegmentation.__init__(s   !'!9!9+F3f&8&8f>S>ST\\&"4"4f6H6HI mm5IaIaCb$ca\&%9$cd .v6(0!yy););V=N=NQR=RS ++v/@/@@&BSBSW]WhWhBhi"("5"5++++.  %F@P@PQ . 6;L;L0MNYYv'9'96;M;MN '%ds>Hrrr5r6rr4c|j|||||}|jjD]'\}}|jD]\} } || vs | |z} )|S)N)rrr5r6r)r(r\r) r.rrr5r6rrrr=loss_keyrs r1 get_loss_dictz-VideomtForUniversalSegmentation.get_loss_dictEs|(,~~!5!5#%"7 (6(  ++113 #KC"+//"3 #$(?FND # # r2rc4t|jSrs)rr)r.rs r1get_lossz(VideomtForUniversalSegmentation.get_loss]s9##%&&r2N patch_offsetsrc Jd|vr td| td|jdk7r td|| td|j\}}}} } |j||z|| | } |j | } |j |j jz } |jd| D] }|| } | j||| jd| jd } g}g}g}d}t|D]S}| dd|f}|2|jjdddddfj|d d }nK|j||jjdddddfj|j z}t#j$|j|j |fd }|j| dD] }||} |j'|}|j)|\}}|j+||j+||j+||ddd|j j,ddf}Vt/dt#j$|d t#j$|d t#j$|d  S)a pixel_values_videos (`torch.Tensor`, *optional*): Video inputs of shape `(batch_size, num_frames, num_channels, height, width)`. mask_labels (`list[torch.Tensor]`, *optional*): Not supported for 5D video inputs. class_labels (`list[torch.LongTensor]`, *optional*): Not supported for 5D video inputs. patch_offsets (`list[torch.Tensor]`, *optional*): Unused for video inputs and only kept for modular compatibility. r3zAUse `pixel_values_videos` with `VideomtForUniversalSegmentation`.Nz'You have to specify pixel_values_videosrczyVideomtForUniversalSegmentation only supports 5D video inputs of shape (batch_size, num_frames, channels, height, width).zTraining with 5D video inputs is not supported in `VideomtForUniversalSegmentation`. Flatten frames and use `EomtForUniversalSegmentation` instead.rr9rNrer)rrrr )r;rgr:rhr@rr/r rrr;rr=r_rr<rdrFrlrpredictrBrr)r.rr5r6r$rrmrnr%rorpflat_pixel_valuesrquery_start_idx layer_moduleall_masks_queries_logitsall_class_queries_logitsall_last_hidden_statespropagated_query frame_idxframe_hidden_states query_tokenssequence_outputrrs r1rAz'VideomtForUniversalSegmentation.forward`s* V #`a a  &FG G  # #q (E   "l&>Q  ?R>W>W; J fe/77 Z8OQ]_eglm(9: 004;;3I3II KK(89 8L(7M 8&**:z=CVCVWXCY[h[n[nop[qr #% #% !#z* TI"/9 "= '#zz00q!<CCJPRTVW #112BCdjjFWFWX\^_abXbFcFfFf'..G  #())\__=P=W=W-XZm,ntu"v  $ O,< = H &23F&G# H#nn-@AO9=o9V 6 "6 $ + +,@ A $ + +,@ A " ) )/ :216O 8O8O6OQR3RS ) T,5!&+C!K!&+C!K#ii(>AF   r2c.|jjSrs)r@rVrs r1get_input_embeddingsz4VideomtForUniversalSegmentation.get_input_embeddingss///r2rc|ddd|jjddf}|j|}|dd|jj|jjzdddf}|j dd}|j |jddg|j}|j|}|j|}tjd||}||fS)Nrr9rrNzbqc, bchw -> bqhw) r/rrr@rZr?rhr:rrrrFeinsum)r.rr0 class_logits prefix_tokens mask_logitss r1r&z'VideomtForUniversalSegmentation.predictsa!:4;;#:#:!:A=> ++L9 q$++"9"9DOO<]<]"]"_abbc %//15 - --m.A.A!.DbZ4>>Z ~~l3 **=9 ll#6 mT L((r2)NNNN)rBrCrDrrr"rrrr!r#rrrrFrQrrrrAr3r&rGrHs@r1rr sF ,O}:$%    $CK0  c6k 0'$sF{"3'' 48152637 O "\\D0O %,,'$.O 5<<(4/ O ELL)D0 O +, O  /O  O b0)ell)r2r)r)rF)F)Mcollections.abcr(rr dataclassesrnumpyrrFtorch.nn.functionalrrrrrr activationsr file_utilsr r r modeling_layersr modeling_utilsrrprocessing_utilsrutilsrrr utils.genericrutils.output_capturingrconfiguration_videomtrscipy.optimizer accelerateraccelerate.utilsrrrrJrvrrrrrirrrrrrrrr-r/rzrUrYr[rrrrrrr__all__r r2r1rKs+* $! &!LL9F&PP7504'' RYY F, ,^&(bii(0% II% <<% % << % LL4' %  %%.8)ryy8)v U\\ e T V[VbVb  %bii %(ryy("'-'T+ +  7K7  7<LQLL5:\\ \\@6,u||X]XdXd8gbiigTff<u||U\\VY^c^j^j(u"))up 1._1.1.h 2  bii" `)&<`)  `)F $%F Gr2