U 9%e@sdZddlZddlZddlmZddlmZmZm Z m Z m Z m Z ddl Z ddlZ ddl mZmZddlmZddlmZmZdd lmZdd lmZmZdd lmZmZmZmZm Z m!Z!m"Z"m#Z#m$Z$d d l%m&Z&e rddl'm(Z(e!rddl)m*Z*e"+e,Z-dZ.dZ/d ddgZ0dgZ1eGdddeZ2Gdddej3Z4Gdddej3Z5Gdddej3Z6Gdddej3Z7Gdddej3Z8Gd d!d!ej3Z9Gd"d#d#ej3Z:Gd$d%d%ej3Z;Gd&d'd'ej3ZGd,d-d-eZ?d.Z@d/ZAed0e@Gd1d2d2e?ZBGd3d4d4ej3ZCed5e@Gd6d7d7e?ZDd8d9ZEdVeFeFd<d=d>ZGGd?d@d@ej3ZHGdAdBdBej3ZIGdCdDdDej3ZJeedEdFdGZKeedHdIdJZLdKdLZMdMdNZNdOdPZOGdQdRdRePZQeedSdTdUZRdS)Wz PyTorch YOLOS model.N) dataclass)DictListOptionalSetTupleUnion)Tensornn)ACT2FN)BaseModelOutputBaseModelOutputWithPooling)PreTrainedModel) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardis_scipy_availableis_vision_availableloggingreplace_return_docstringsrequires_backends) YolosConfiglinear_sum_assignment)center_to_corners_formatrzhustvl/yolos-smalliI ic@seZdZUdZdZeejed<dZ ee ed<dZ ejed<dZ ejed<dZ eee ed<dZeejed<dZeeejed <dZeeejed <dS) YolosObjectDetectionOutputaG Output type of [`YolosForObjectDetection`]. Args: 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 auxilary 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. 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, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nloss loss_dictlogits pred_boxesauxiliary_outputslast_hidden_state hidden_states attentions)__name__ __module__ __qualname____doc__r!rtorch FloatTensor__annotations__r"rr#r$r%rr&r'rr(r0r0g/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/models/yolos/modeling_yolos.pyr Ds !r cs<eZdZdZeddfdd ZejejdddZZ S) YolosEmbeddingszT Construct the CLS token, detection tokens, position and patch embeddings. Nconfigreturncstttdd|j|_ttd|j|j|_ t ||_ |j j }ttd||jd|j|_ t|j|_t||_||_dSNr)super__init__r Parameterr-zeros hidden_size cls_tokennum_detection_tokensdetection_tokensYolosPatchEmbeddingspatch_embeddings num_patchesposition_embeddingsDropouthidden_dropout_probdropout$InterpolateInitialPositionEmbeddings interpolationr4)selfr4rA __class__r0r1r8ws   zYolosEmbeddings.__init__ pixel_valuesr5c Cs|j\}}}}||}|\}}}|j|dd} |j|dd} tj| || fdd}||j ||f} || }| |}|S)Nrdim) shaper@sizer<expandr>r-catrGrBrE) rHrL batch_size num_channelsheightwidth embeddingsseq_len_Z cls_tokensr>rBr0r0r1forwards  zYolosEmbeddings.forward r)r*r+r,rr8r-r r[ __classcell__r0r0rIr1r2qsr2cs4eZdZddfdd ZdejdddZZS) rFNr5cst||_dSNr7r8r4rHr4rIr0r1r8s z-InterpolateInitialPositionEmbeddings.__init__i i@cCs|dddddf}|dddf}|dd|jj dddf}|ddd|jj ddf}|dd}|j\}}}|jjd|jj|jjd|jj} } |||| | }|\} } | |jj| |jj} }tjj || |fddd}| ddd}t j |||fdd}|S)NrrbicubicFrQmodeZ align_cornersrN) r4r= transposerP image_size patch_sizeviewr functional interpolateflattenr-rS)rH pos_embedimg_size cls_pos_embed det_pos_embedpatch_pos_embedrTr;rY patch_height patch_widthrVrWZnew_patch_heigthnew_patch_widthscale_pos_embedr0r0r1r[s*    z,InterpolateInitialPositionEmbeddings.forward)rbr)r*r+r8r-r r[r]r0r0rIr1rFsrFcs4eZdZddfdd ZdejdddZZS) InterpolateMidPositionEmbeddingsNr^cst||_dSr_r`rarIr0r1r8s z)InterpolateMidPositionEmbeddings.__init__rbcCsH|dddddddf}|dddf}|dddd|jj dddf}|ddddd|jj ddf}|dd}|j\}}}} |jjd|jj|jjd|jj} } ||||| | }|\} } | |jj| |jj}}tjj |||fddd}| ddd |||||}t j |||fdd}|S) Nrrrcr rdFrerN)r4r=rgrPrhrirjr rkrlrm contiguousr-rS)rHrnrorprqrrdepthrTr;rYrsrtrVrWZnew_patch_heightrurvr0r0r1r[s<&&  z(InterpolateMidPositionEmbeddings.forward)rbrwr0r0rIr1rxsrxcs4eZdZdZfddZejejdddZZS)r?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. cst|j|j}}|j|j}}t|tjj r8|n||f}t|tjj rR|n||f}|d|d|d|d}||_||_||_||_ t j ||||d|_ dS)Nrr)Z kernel_sizeZstride)r7r8rhrirUr; isinstance collectionsabcIterablerAr Conv2d projection)rHr4rhrirUr;rArIr0r1r8s  zYolosPatchEmbeddings.__init__rKcCs<|j\}}}}||jkr td||ddd}|S)NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.rcr)rPrU ValueErrorrrmrg)rHrLrTrUrVrWrXr0r0r1r[s zYolosPatchEmbeddings.forward) r)r*r+r,r8r-r r[r]r0r0rIr1r?s r?csleZdZeddfdd ZejejdddZd eeje e e ejejfe ejfd d d Z Z S) YolosSelfAttentionNr3cst|j|jdkr@t|ds@td|jfd|jd|j|_t|j|j|_|j|j|_t j |j|j|j d|_ t j |j|j|j d|_ t j |j|j|j d|_t |j|_dS)NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .)bias)r7r8r;num_attention_headshasattrrintattention_head_size all_head_sizer LinearZqkv_biasquerykeyvaluerCZattention_probs_dropout_probrErarIr0r1r8s zYolosSelfAttention.__init__)xr5cCs6|dd|j|jf}||}|ddddS)NrMrrcrr )rQrrrjpermute)rHrZ new_x_shaper0r0r1transpose_for_scoress z'YolosSelfAttention.transpose_for_scoresF) head_maskoutput_attentionsr5c Cs||}|||}|||}||}t||dd}|t|j }t j j |dd} | | } |dk r| |} t| |} | dddd} | dd|jf} | | } |r| | fn| f} | S)NrMrNrrcrr )rrrrr-matmulrgmathsqrtrr rksoftmaxrErryrQrrj) rHr'rrZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr0r0r1r[s      zYolosSelfAttention.forward)NF)r)r*r+rr8r-r rrboolrrr[r]r0r0rIr1rsrcs@eZdZdZeddfdd ZejejejdddZZ S) YolosSelfOutputz 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. Nr3cs.tt|j|j|_t|j|_dSr_) r7r8r rr;denserCrDrErarIr0r1r8?s zYolosSelfOutput.__init__r' input_tensorr5cCs||}||}|Sr_rrErHr'rr0r0r1r[Ds  zYolosSelfOutput.forwardr\r0r0rIr1r9srcspeZdZeddfdd ZeeddddZd ej e ej e e e ej ej fe ej fd d d ZZS) YolosAttentionNr3cs*tt||_t||_t|_dSr_)r7r8r attentionroutputset pruned_headsrarIr0r1r8Ms   zYolosAttention.__init__)headsr5cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)NrrrN)lenrrrrrrrrrrrrunion)rHrindexr0r0r1 prune_headsSs zYolosAttention.prune_headsFr'rrr5cCs4||||}||d|}|f|dd}|SNrr)rr)rHr'rrZ self_outputsattention_outputrr0r0r1r[eszYolosAttention.forward)NF)r)r*r+rr8rrrr-r rrrrr[r]r0r0rIr1rLsrcs8eZdZeddfdd ZejejdddZZS)YolosIntermediateNr3csBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr_) r7r8r rr;intermediate_sizerr{Z hidden_actstrr intermediate_act_fnrarIr0r1r8us   zYolosIntermediate.__init__)r'r5cCs||}||}|Sr_)rr)rHr'r0r0r1r[}s  zYolosIntermediate.forward r)r*r+rr8r-r r[r]r0r0rIr1rtsrcs<eZdZeddfdd ZejejejdddZZS) YolosOutputNr3cs.tt|j|j|_t|j|_ dSr_) r7r8r rrr;rrCrDrErarIr0r1r8s zYolosOutput.__init__rcCs ||}||}||}|Sr_rrr0r0r1r[s  zYolosOutput.forwardrr0r0rIr1rsrcs`eZdZdZeddfdd Zd ejeeje e e ejejfe ejfddd Z Z S) YolosLayerz?This corresponds to the Block class in the timm implementation.Nr3csbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dS)NrZeps)r7r8Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr LayerNormr;layer_norm_epslayernorm_beforelayernorm_afterrarIr0r1r8s    zYolosLayer.__init__FrcCs`|j||||d}|d}|dd}||}||}||}|||}|f|}|S)Nrrr)rrrrr)rHr'rrZself_attention_outputsrrZ layer_outputr0r0r1r[s     zYolosLayer.forward)NF)r)r*r+r,rr8r-r rrrrr[r]r0r0rIr1rs rc sNeZdZeddfdd Zd ejeejeeee e e fddd Z Z S) YolosEncoderNr3cst|_tfddtjD|_d|_dj dj dj dj }j rt tjdd|jnd|_j rtnd|_dS)Ncsg|] }tqSr0)r).0rZr4r0r1 sz)YolosEncoder.__init__..Frrrc)r7r8r4r ModuleListrangenum_hidden_layerslayergradient_checkpointingrhrir=use_mid_position_embeddingsr9r-r:r;mid_position_embeddingsrxrG)rHr4Z seq_lengthrIrr1r8s$  &   zYolosEncoder.__init__FT)r'rroutput_hidden_states return_dictr5cs|rdnd}rdnd} |jjr2||j||f} t|jD]\} } |rR||f}|dk rb|| nd} |jr|jrfdd}tj j || || }n | || }|d}|jjr| |jj dkr|| | }r<| |df} q<|r||f}|st dd||| fDSt ||| dS) Nr0csfdd}|S)Ncs|fSr_r0)inputs)modulerr0r1custom_forwardszKYolosEncoder.forward..create_custom_forward..custom_forwardr0)rrr)rr1create_custom_forwardsz3YolosEncoder.forward..create_custom_forwardrrcss|]}|dk r|VqdSr_r0rvr0r0r1 sz'YolosEncoder.forward..)r&r'r()r4rrGr enumeraterrZtrainingr-utils checkpointrtupler )rHr'rVrWrrrrZall_hidden_statesZall_self_attentionsZ$interpolated_mid_position_embeddingsiZ layer_moduleZlayer_head_maskrZ layer_outputsr0rr1r[s>        zYolosEncoder.forward)NFFT)r)r*r+rr8r-r rrrrr r[r]r0r0rIr1rs rc@sTeZdZdZeZdZdZdZe e j e j e j fddddZd eedd d d ZdS)YolosPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. vitrLTN)rr5cCsjt|tjtjfr@|jjjd|jjd|j dk rf|j j n&t|tj rf|j j |jj ddS)zInitialize the weightsg)meanZstdNg?) r{r rrweightdataZnormal_r4Zinitializer_rangerZzero_rZfill_)rHrr0r0r1 _init_weightss   z"YolosPreTrainedModel._init_weightsF)rrr5cCst|tr||_dSr_)r{rr)rHrrr0r0r1_set_gradient_checkpointing*s z0YolosPreTrainedModel._set_gradient_checkpointing)F)r)r*r+r,r config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrr rrrrrrrr0r0r0r1rs raG This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`YolosConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. aM Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`YolosImageProcessor.__call__`] for details. head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. z_The bare YOLOS Model transformer outputting raw hidden-states without any specific head on top.c seZdZdeedfdd ZedddZee e e fdd d d Z e e eeeed ed deejeejeeeeeeeeefdddZZS) YolosModelT)r4add_pooling_layercsXt|||_t||_t||_tj|j |j d|_ |rFt |nd|_ |dS)Nr)r7r8r4r2rXrencoderr rr;r layernorm YolosPoolerpooler post_init)rHr4rrIr0r1r8Vs   zYolosModel.__init__r^cCs|jjSr_)rXr@rHr0r0r1get_input_embeddingscszYolosModel.get_input_embeddingsN)heads_to_pruner5cCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. Args: heads_to_prune (`dict` of {layer_num: list of heads to prune in this layer}): See base class `PreTrainedModel`. N)itemsrrrr)rHrrrr0r0r1 _prune_headsfszYolosModel._prune_headsZvision)r output_typerZmodalityZexpected_output)rLrrrrr5c Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|dkrLtd|||jj}||}|j||j d|j d||||d}|d}| |}|j dk r| |nd} |s| dk r|| fn|f} | |ddSt || |j |jdS)Nz You have to specify pixel_valuesrrM)rVrWrrrrrr)r&Z pooler_outputr'r()r4rruse_return_dictrZ get_head_maskrrXrrPrrrr'r() rHrLrrrrZembedding_outputZencoder_outputssequence_output pooled_outputZ head_outputsr0r0r1r[qs:   zYolosModel.forward)T)NNNNN)r)r*r+rrr8r?rrrrrrYOLOS_INPUTS_DOCSTRINGr_CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPErr-r rrr[r]r0r0rIr1rQs0    rcs*eZdZedfdd ZddZZS)rrcs*tt|j|j|_t|_dSr_)r7r8r rr;rZTanh activationrarIr0r1r8s zYolosPooler.__init__cCs(|dddf}||}||}|S)Nr)rr)rHr'Zfirst_token_tensorrr0r0r1r[s  zYolosPooler.forward)r)r*r+rr8r[r]r0r0rIr1rsrzy YOLOS Model (consisting of a ViT encoder) with object detection heads on top, for tasks such as COCO detection. c s~eZdZedfdd ZejjddZe e e e e dd ejeeeeeeeeeeee fdd d ZZS) YolosForObjectDetectionrcsXt|t|dd|_t|j|j|jddd|_t|j|jddd|_| dS)NF)rrr ) input_dim hidden_dim output_dim num_layers) r7r8rrYolosMLPPredictionHeadr; num_labelsclass_labels_classifierbbox_predictorrrarIr0r1r8s z YolosForObjectDetection.__init__cCs$ddt|dd|ddDS)NcSsg|]\}}||dqS))r#r$r0)rabr0r0r1rsz9YolosForObjectDetection._set_aux_loss..rM)zip)rH outputs_class outputs_coordr0r0r1 _set_aux_losssz%YolosForObjectDetection._set_aux_loss)rrN)rLlabelsrrrr5c s|dk r |n|jj}|j||||d}|d}|dd|jj dddf}||}||} d\} } |dk rt|jj|jj |jj d} dddg} t | |jj |jj | d }||ji}||d <| |d <|jjr|r|jn|d }||}||}|||} | |d <|||d|jjd|jjd<|jjri}t|jjdD]"|fddDq`|tfddD} |s| dk r|| f| |}n || f|}| dk r| f|S|St| || | |j|j|jdS)a 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)`. Returns: Examples: ```python >>> from transformers import AutoImageProcessor, AutoModelForObjectDetection >>> import torch >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> 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 COCO API >>> 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.994 at location [46.96, 72.61, 181.02, 119.73] Detected remote with confidence 0.975 at location [340.66, 79.19, 372.59, 192.65] Detected cat with confidence 0.984 at location [12.27, 54.25, 319.42, 470.99] Detected remote with confidence 0.922 at location [41.66, 71.96, 178.7, 120.33] Detected cat with confidence 0.914 at location [342.34, 21.48, 638.64, 372.46] ```N)rrrr)NNN class_cost bbox_cost giou_costrboxes cardinality)matcher num_classeseos_coeflossesr#r$rr%r)loss_ce loss_bbox loss_gioucs i|]\}}|d|qSrZr0rkrrr0r1 Gs z3YolosForObjectDetection.forward..c3s&|]}|kr||VqdSr_r0)rr)r" weight_dictr0r1rIsz2YolosForObjectDetection.forward..)r!r"r#r$r%r&r'r()r4rrr=rrsigmoidYolosHungarianMatcherrrr YolosLossrZeos_coefficienttodeviceZauxiliary_lossZintermediate_hidden_statesrZbbox_loss_coefficientZgiou_loss_coefficientrZdecoder_layersupdatersumkeysr r&r'r()rHrLrrrrrrr#r$r!r%rr  criterionZ outputs_lossrrrZaux_weight_dictrr0)rr"rr1r[sv7                zYolosForObjectDetection.forward)NNNN)r)r*r+rr8r-ZjitZunusedrrrrr rr.rrrrrrr[r]r0r0rIr1rs"    rcCsX|}|d}d||d}|d|d}d|d|d}||S)a Compute the DICE loss, similar to generalized IOU for masks Args: inputs: A float tensor of arbitrary shape. The predictions for each example. targets: A float tensor with the same shape as inputs. Stores the binary classification label for each element in inputs (0 for the negative class and 1 for the positive class). rrcrM)rrmr)rtargets num_boxes numerator denominatorr!r0r0r1 dice_loss_s  r"?rc)alphagammac Cs||}tjj||dd}||d|d|}|d||}|dkrj||d|d|} | |}|d|S)a Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. Args: inputs (`torch.FloatTensor` of arbitrary shape): The predictions for each example. targets (`torch.FloatTensor` with the same shape as `inputs`) A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class and 1 for the positive class). alpha (`float`, *optional*, defaults to `0.25`): Optional weighting factor in the range (0,1) to balance positive vs. negative examples. gamma (`int`, *optional*, defaults to `2`): Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples. Returns: Loss tensor noneZ reductionrr)rr rkZ binary_cross_entropy_with_logitsrr) rrrr$r%ZprobZce_lossZp_tr!Zalpha_tr0r0r1sigmoid_focal_lossssr(csheZdZdZfddZddZeddZdd Z d d Z d d Z ddZ ddZ ddZZS)ra This class computes the losses for YolosForObjectDetection/YolosForSegmentation. The process happens in two steps: 1) we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched ground-truth / prediction (supervise class and box). A note on the `num_classes` argument (copied from original repo in detr.py): "the naming of the `num_classes` parameter of the criterion is somewhat misleading. It indeed corresponds to `max_obj_id` + 1, where `max_obj_id` is the maximum id for a class in your dataset. For example, COCO has a `max_obj_id` of 90, so we pass `num_classes` to be 91. As another example, for a dataset that has a single class with `id` 1, you should pass `num_classes` to be 2 (`max_obj_id` + 1). For more details on this, check the following discussion https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223" Args: matcher (`YolosHungarianMatcher`): Module able to compute a matching between targets and proposals. num_classes (`int`): Number of object categories, omitting the special no-object category. eos_coef (`float`): Relative classification weight applied to the no-object category. losses (`List[str]`): List of all the losses to be applied. See `get_loss` for a list of all available losses. csLt||_||_||_||_t|jd}|j|d<|d|dS)NrrM empty_weight) r7r8rr r r r-onesZregister_buffer)rHrr r r r)rIr0r1r8s  zYolosLoss.__init__c Csd|krtd|d}||}tddt||D}tj|jdd|jtj|j d}|||<t j | dd||j} d | i} | S) z Classification loss (NLL) targets dicts must contain the key "class_labels" containing a tensor of dim [nb_target_boxes] r#z#No logits were found in the outputscSs g|]\}\}}|d|qS class_labelsr0)rtrZJr0r0r1rs z)YolosLoss.loss_labels..Nrcdtyperrr )KeyError_get_source_permutation_idxr-rSrfullrPr int64rr rkZ cross_entropyrgr)) rHrrindicesrZ source_logitsidxZtarget_classes_oZtarget_classesr r r0r0r1 loss_labelss  zYolosLoss.loss_labelsc Csf|d}|j}tjdd|D|d}|d|jddkd}tj| | } d| i} | S)z Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes. This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients. r#cSsg|]}t|dqSr+rrr0r0r1rsz.YolosLoss.loss_cardinality..)rrMrZcardinality_error) rr- as_tensorZargmaxrPrr rkl1_lossfloat) rHrrr5rr#rZtarget_lengthsZ card_predZcard_errr r0r0r1loss_cardinalityszYolosLoss.loss_cardinalityc Csd|krtd||}|d|}tjddt||Ddd}tjj||dd}i} ||| d <d t t t |t |} | || d <| S) a< Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss. Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size. r$z#No predicted boxes found in outputscSs g|]\}\}}|d|qSrr0)rr-rZrr0r0r1rs z(YolosLoss.loss_boxes..rrNr&r'r rr) r1r2r-rSrr rkr:rZdiaggeneralized_box_iour) rHrrr5rr6Z source_boxesZ target_boxesr r rr0r0r1 loss_boxess  zYolosLoss.loss_boxesc Csd|krtd||}||}|d}||}dd|D}t|\} } | |} | |} tjj|dddf| j ddddd }|ddd f d }| d } | |j } t || |t || |d } | S) z Compute the losses related to the masks: the focal loss and the dice loss. Targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]. Z pred_masksz#No predicted masks found in outputscSsg|] }|dqS)masksr0rr-r0r0r1rsz(YolosLoss.loss_masks..NrZbilinearFrerr)Z loss_maskZ loss_dice)r1r2_get_target_permutation_idxnested_tensor_from_tensor_list decomposerr rkrlrPrmrjr(r") rHrrr5r source_idx target_idxZ source_masksr@Z target_masksZvalidr r0r0r1 loss_maskss.        zYolosLoss.loss_maskscCs4tddt|D}tdd|D}||fS)NcSs g|]\}\}}t||qSr0r-Z full_like)rrsourcerZr0r0r1rs z9YolosLoss._get_source_permutation_idx..cSsg|] \}}|qSr0r0)rrIrZr0r0r1rsr-rSr)rHr5 batch_idxrEr0r0r1r2sz%YolosLoss._get_source_permutation_idxcCs4tddt|D}tdd|D}||fS)NcSs g|]\}\}}t||qSr0rH)rrrZtargetr0r0r1rs z9YolosLoss._get_target_permutation_idx..cSsg|] \}}|qSr0r0)rrZrLr0r0r1rsrJ)rHr5rKrFr0r0r1rBsz%YolosLoss._get_target_permutation_idxcCs@|j|j|j|jd}||kr.td|d||||||S)N)rrrr@zLoss z not supported)r7r<r?rGr)rHr!rrr5rZloss_mapr0r0r1get_loss!szYolosLoss.get_lossc sdd|D}|||}tdd|D}tj|gtjtt|j d}tj |dd }i}|j D]}| ||||||qnd|krt|dD]\\}|||}|j D]@}|d krq||||||} fd d| D} | | qq|S) a This performs the loss computation. Args: outputs (`dict`, *optional*): Dictionary of tensors, see the output specification of the model for the format. targets (`List[dict]`, *optional*): List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the losses applied, see each loss' doc. cSsi|]\}}|dkr||qS)r%r0rr0r0r1r7sz%YolosLoss.forward..css|]}t|dVqdS)r,Nr8rAr0r0r1r=sz$YolosLoss.forward..r/rminr%r@cs i|]\}}|d|qSrr0rrr0r1rSs )rrrr-r9r;nextitervaluesrclampitemr rrMr) rHrrZoutputs_without_auxr5rr r!r%Zl_dictr0rr1r[,s$  "   zYolosLoss.forward)r)r*r+r,r8r7r-no_gradr<r?rGr2rBrMr[r]r0r0rIr1rs  ! rcs(eZdZdZfddZddZZS)ra 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. Copied from https://github.com/facebookresearch/detr/blob/master/models/detr.py csJt||_|g|d}tddt|g|||gD|_dS)Nrcss|]\}}t||VqdSr_)r r)rnrr0r0r1rgsz2YolosMLPPredictionHead.__init__..)r7r8rr rrlayers)rHrrrrhrIr0r1r8cs zYolosMLPPredictionHead.__init__cCs>t|jD].\}}||jdkr0tj||n||}q |Sr6)rrWrr rkZrelu)rHrrrr0r0r1r[is(zYolosMLPPredictionHead.forward)r)r*r+r,r8r[r]r0r0rIr1rZs rcs<eZdZdZdeeedfdd ZeddZZ S) ra This class computes an assignment between the targets and the predictions of the network. For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more predictions than targets. 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). Args: class_cost: The relative weight of the classification error in the matching cost. bbox_cost: The relative weight of the L1 error of the bounding box coordinates in the matching cost. giou_cost: The relative weight of the giou loss of the bounding box in the matching cost. rrcsLtt|dg||_||_||_|dkrH|dkrH|dkrHtddS)NZscipyrz#All costs of the Matcher can't be 0)r7r8rrrrr)rHrrrrIr0r1r8s  zYolosHungarianMatcher.__init__cCs|djdd\}}|dddd}|ddd}tdd |D}td d |D}|dd|f } tj||dd } tt|t| } |j| |j | |j | } | ||d } d d |D} d d t | | dD}dd |DS)a Args: outputs (`dict`): A dictionary that contains at least these entries: * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates. targets (`List[dict]`): A list of targets (len(targets) = batch_size), where each target is a dict containing: * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth objects in the target) containing the class labels * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates. Returns: `List[Tuple]`: 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 targets (in order) For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes) r#NrcrrrMr$cSsg|] }|dqSr+r0rr0r0r1rsz1YolosHungarianMatcher.forward..cSsg|] }|dqSr=r0rr0r0r1rs)pcSsg|]}t|dqSr=r8rr0r0r1rscSsg|]\}}t||qSr0r)rrcr0r0r1rscSs0g|](\}}tj|tjdtj|tjdfqS))r0)r-r9r4)rrjr0r0r1rs)rPrmrr-rSZcdistr>rrrrrjcpursplit)rHrrrTZ num_queriesZout_probZout_bboxZ target_idsZ target_bboxrrrZ cost_matrixsizesr5r0r0r1r[szYolosHungarianMatcher.forward)rrr) r)r*r+r,r;r8r-rUr[r]r0r0rIr1rps r)r-r5cCsH|r&|jtjtjfkr|S|S|jtjtjfkr<|S|SdSr_) Zis_floating_pointr0r-Zfloat32Zfloat64r;Zint32r4r)r-r0r0r1_upcastsr_)rr5cCsHt|}|dddf|dddf|dddf|dddfS)a Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates. Args: boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`): Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1 < x2` and `0 <= y1 < y2`. Returns: `torch.FloatTensor`: a tensor containing the area for each box. Nrcrr r)r_r=r0r0r1box_areas r`c Cst|}t|}t|dddddf|ddddf}t|dddddf|ddddf}||jdd}|dddddf|dddddf}|dddf||}||} | |fS)NrcrrNr)r`r-maxrOrS) boxes1boxes2Zarea1Zarea2Zleft_topZ right_bottom width_heightZinterriour0r0r1box_ious..,rfcCs*|ddddf|ddddfks:td||ddddf|ddddfksttd|t||\}}t|dddddf|ddddf}t|dddddf|ddddf}||jdd}|dddddf|dddddf}||||S)z Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format. Returns: `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2) Nrczs ,,..,r>cCsB|d}|ddD](}t|D]\}}t|||||<q q|Sr)rra)Zthe_listZmaxesZsublistrrTr0r0r1 _max_by_axiss rjc@s6eZdZeedddZddZddZdd Zd S) NestedTensor)maskcCs||_||_dSr_tensorsrl)rHrnrlr0r0r1r8szNestedTensor.__init__cCs4|j|}|j}|dk r&||}nd}t||Sr_)rnrrlrk)rHrZ cast_tensorrlZ cast_maskr0r0r1rs   zNestedTensor.tocCs |j|jfSr_rmrr0r0r1rDszNestedTensor.decomposecCs t|jSr_)rrnrr0r0r1__repr__szNestedTensor.__repr__N) r)r*r+rr r8rrDror0r0r0r1rk s rk) tensor_listcCs|djdkrtdd|D}t|g|}|\}}}}|dj}|dj}tj|||d} tj|||ftj|d} t || | D]\\} } } | d| j dd| j dd| j df | d| d| j dd| j df<qnt d t | | S) Nrr cSsg|]}t|jqSr0)listrP)rimgr0r0r1r%sz2nested_tensor_from_tensor_list..r/rrcFz(Only 3-dimensional tensors are supported)ndimrjrr0rr-r:r*rrrPZcopy_rrk)rpmax_sizeZ batch_shaperTrUrVrWr0rZtensorrlrrZpad_imgmr0r0r1rC#s   2$rC)r#rc)Sr,collections.abcr|r dataclassesrtypingrrrrrrr-Ztorch.utils.checkpointr r Z activationsr Zmodeling_outputsr rZmodeling_utilsrZ pytorch_utilsrrrrrrrrrrrrZconfiguration_yolosrZscipy.optimizerZtransformers.image_transformsrZ get_loggerr)loggerrrrZ#YOLOS_PRETRAINED_MODEL_ARCHIVE_LISTr Moduler2rFrxr?rrrrrrrrZYOLOS_START_DOCSTRINGrrrrr"r;r(rrrr_r`rfr>rjobjectrkrCr0r0r0r1s    ,     ,+!"=(*T W HO