o h8@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 rqddl'm(Z(e!rzddl)m*Z*e"+e,Z-dZ.dZ/gdZ0dgZ1eGdddeZ2Gdddej3Z4Gdddej3Z5Gdddej3Z6Gdddej3Z7Gdddej3Z8Gdd d ej3Z9Gd!d"d"ej3Z:Gd#d$d$ej3Z;Gd%d&d&ej3ZGd+d,d,eZ?d-Z@d.ZAed/e@Gd0d1d1e?ZBGd2d3d3ej3ZCed4e@Gd5d6d6e?ZDd7d8ZEdWd;eFdZGGd?d@d@ej3ZHGdAdBdBej3ZIGdCdDdDej3ZJdEedFefdGdHZKdIedFefdJdKZLdLdMZMdNdOZNdPdQZOGdRdSdSePZQdTeefdUdVZRdS)Xz 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-small)riI 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(r0r0^/var/www/html/ai/venv/lib/python3.10/site-packages/transformers/models/yolos/modeling_yolos.pyr Ds !r cs@eZdZdZdeddffdd Zdejdejfdd ZZ S) YolosEmbeddingszT Construct the CLS token, detection tokens, position and patch embeddings. configreturnNcstttdd|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 interpolationr3)selfr3r@ __class__r0r1r7ws    zYolosEmbeddings.__init__ pixel_valuesc Cs|j\}}}}||}|\}}}|j|dd} |j|dd} tj| || fdd}||j ||f} || }| |}|S)Nrdim) shaper?sizer;expandr=r-catrFrArD) rGrJ batch_size num_channelsheightwidth embeddingsseq_len_ cls_tokensr=rAr0r0r1forwards  zYolosEmbeddings.forward r)r*r+r,rr7r-r rZ __classcell__r0r0rHr1r2qsr2c0eZdZdfdd Zd dejfddZZS) rEr4Nct||_dSNr6r7r3rGr3rHr0r1r7  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)NrrbicubicFrOmode align_cornersrL) r3r< transposerN image_size patch_sizeviewr functional interpolateflattenr-rQ)rG pos_embedimg_size cls_pos_embed det_pos_embedpatch_pos_embedrRr:rW patch_height patch_widthrTrUnew_patch_heigthnew_patch_widthscale_pos_embedr0r0r1rZs$     z,InterpolateInitialPositionEmbeddings.forwardr4Nrcr)r*r+r7r-r rZr\r0r0rHr1rErEcr]) InterpolateMidPositionEmbeddingsr4Ncr^r_r`rarHr0r1r7rbz)InterpolateMidPositionEmbeddings.__init__rccCsH|dddddddf}|dddf}|dddd|jj dddf}|ddddd|jj ddf}|dd}|j\}}}} |jjd|jj|jjd|jj} } ||||| | }|\} } | |jj| |jj}}tjj |||fddd}| ddd |||||}t j |||fdd}|S) Nrrrdr reFrfrL)r3r<rirNrjrkrlr rmrnro contiguousr-rQ)rGrprqrrrsrtdepthrRr:rWrurvrTrUnew_patch_heightrxryr0r0r1rZs,&&  z(InterpolateMidPositionEmbeddings.forwardrzr{r|r0r0rHr1r~r}r~cs6eZdZdZfddZdejdejfddZZS)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 r|n||f}t|tjj r)|n||f}|d|d|d|d}||_||_||_||_ t j ||||d|_ dS)Nrr) kernel_sizestride)r6r7rjrkrSr: isinstance collectionsabcIterabler@r Conv2d projection)rGr3rjrkrSr:r@rHr0r1r7s  zYolosPatchEmbeddings.__init__rJr4cCs<|j\}}}}||jkrtd||ddd}|S)NzeMake sure that the channel dimension of the pixel values match with the one set in the configuration.rdr)rNrS ValueErrorrrori)rGrJrRrSrTrUrVr0r0r1rZs zYolosPatchEmbeddings.forward) r)r*r+r,r7r-r rZr\r0r0rHr1r>s r>c sveZdZdeddffdd ZdejdejfddZ dd eejd e de e ejejfe ejffd d Z Z S)YolosSelfAttentionr3r4Ncst|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)Nrembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .)bias)r6r7r:num_attention_headshasattrrintattention_head_size all_head_sizer Linearqkv_biasquerykeyvaluerBattention_probs_dropout_probrDrarHr0r1r7s  zYolosSelfAttention.__init__xcCs6|dd|j|jf}||}|ddddS)NrKrrdrr )rOrrrlpermute)rGr new_x_shaper0r0r1transpose_for_scoress z'YolosSelfAttention.transpose_for_scoresF head_maskoutput_attentionsc Cs||}|||}|||}||}t||dd}|t|j }t j j |dd} | | } |durA| |} t| |} | dddd} | dd|jf} | | } |rj| | f} | S| f} | S)NrKrLrrdrr )rrrrr-matmulrimathsqrtrr rmsoftmaxrDrrrOrrl) rGr'rrmixed_query_layer key_layer value_layer query_layerattention_scoresattention_probs context_layernew_context_layer_shapeoutputsr0r0r1rZs$      zYolosSelfAttention.forwardNF)r)r*r+rr7r-r rrboolrrrZr\r0r0rHr1rsrcsFeZdZdZdeddffdd Zdejdejdejfd d ZZ 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. r3r4Ncs.tt|j|j|_t|j|_dSr_) r6r7r rr:denserBrCrDrarHr0r1r7? zYolosSelfOutput.__init__r' input_tensorcC||}||}|Sr_rrDrGr'rr0r0r1rZD  zYolosSelfOutput.forwardr[r0r0rHr1r9s$rc s~eZdZdeddffdd ZdeeddfddZ  dd ej d e ej d e de e ej ej fe ej ffd dZZS)YolosAttentionr3r4Ncs*tt||_t||_t|_dSr_)r6r7r attentionroutputset pruned_headsrarHr0r1r7Ms    zYolosAttention.__init__headscCst|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)NrrrL)lenrrrrrrrrrrrrunion)rGrindexr0r0r1 prune_headsSs zYolosAttention.prune_headsFr'rrcCs4||||}||d|}|f|dd}|SNrr)rr)rGr'rr self_outputsattention_outputrr0r0r1rZeszYolosAttention.forwardr)r)r*r+rr7rrrr-r rrrrrZr\r0r0rHr1rLsrcs<eZdZdeddffdd ZdejdejfddZZS) YolosIntermediater3r4NcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr_) r6r7r rr:intermediate_sizerr hidden_actstrr intermediate_act_fnrarHr0r1r7us   zYolosIntermediate.__init__r'cCrr_)rr)rGr'r0r0r1rZ}rzYolosIntermediate.forward r)r*r+rr7r-r rZr\r0r0rHr1rtsrcsBeZdZdeddffdd Zdejdejdejfdd ZZS) YolosOutputr3r4Ncs.tt|j|j|_t|j|_ dSr_) r6r7r rrr:rrBrCrDrarHr0r1r7rzYolosOutput.__init__r'rcCs ||}||}||}|Sr_rrr0r0r1rZs  zYolosOutput.forwardrr0r0rHr1rs$rc sleZdZdZdeddffdd Z  d dejd eejd e de e ejejfe ejffd d Z Z S) YolosLayerz?This corresponds to the Block class in the timm implementation.r3r4Ncsbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dS)Nreps)r6r7chunk_size_feed_forward seq_len_dimrrr intermediaterrr LayerNormr:layer_norm_epslayernorm_beforelayernorm_afterrarHr0r1r7s    zYolosLayer.__init__Fr'rrcCs`|j||||d}|d}|dd}||}||}||}|||}|f|}|S)N)rrr)rrrrr)rGr'rrself_attention_outputsrr layer_outputr0r0r1rZs     zYolosLayer.forwardr)r)r*r+r,rr7r-r rrrrrZr\r0r0rHr1rs rcsbeZdZdeddffdd Z    ddejd eejd ed ed ede e e ff d dZ Z S) YolosEncoderr3r4Ncst|_tfddtjD|_d|_dj dj dj dj }j rAt tjdd|jnd|_j rNt|_dSd|_dS)Ncsg|]}tqSr0)r).0rXr3r0r1 z)YolosEncoder.__init__..Frrrd)r6r7r3r ModuleListrangenum_hidden_layerslayergradient_checkpointingrjrkr<use_mid_position_embeddingsr8r-r9r:mid_position_embeddingsr~rF)rGr3 seq_lengthrHrr1r7s$  &   zYolosEncoder.__init__FTr'rroutput_hidden_states return_dictcCs |rdnd}|r dnd} |jjr||j||f} t|jD]J\} } |r)||f}|dur1|| nd} |jrC|jrC|| j || |}n| || |}|d}|jjr_| |jj dkr_|| | }|rh| |df} q|rp||f}|s~t dd||| fDSt ||| dS)Nr0rrcss|] }|dur|VqdSr_r0rvr0r0r1 z'YolosEncoder.forward..)r&r'r() r3rrFr enumeraterrtraining_gradient_checkpointing_func__call__rtupler )rGr'rTrUrrrrall_hidden_statesall_self_attentions$interpolated_mid_position_embeddingsi layer_modulelayer_head_mask layer_outputsr0r0r1rZs@      zYolosEncoder.forward)NFFT)r)r*r+rr7r-r rrrrr rZr\r0r0rHr1rs& rc@sBeZdZdZeZdZdZdZde e j e j e j fddfdd ZdS) YolosPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. vitrJTmoduler4NcCstt|tjtjfr#|jjjd|jjd|j dur!|j j dSdSt|tj r8|j j |jj ddSdS)zInitialize the weightsg)meanstdNg?) rr rrweightdatanormal_r3initializer_rangerzero_rfill_)rGrr0r0r1 _init_weightss   z"YolosPreTrainedModel._init_weights)r)r*r+r,r config_classbase_model_prefixmain_input_namesupports_gradient_checkpointingrr rrrrr0r0r0r1r s&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.cseZdZddedeffdd ZdefddZd ee e e fdd fd d Z e e eeeed ed ddeejdeejdeedeedeedeeeff ddZZS) YolosModelTr3add_pooling_layercsXt|||_t||_t||_tj|j |j d|_ |r#t |nd|_ |dS)Nr)r6r7r3r2rVrencoderr rr:r layernorm YolosPoolerpooler post_init)rGr3rrHr0r1r7Ls    zYolosModel.__init__r4cCs|jjSr_)rVr?rGr0r0r1get_input_embeddingsYszYolosModel.get_input_embeddingsheads_to_pruneNcCs*|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)itemsr rrr)rGrrrr0r0r1 _prune_heads\szYolosModel._prune_headsvision) checkpoint output_typermodalityexpected_outputrJrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|||jj}||}|j||j d|j d||||d}|d}| |}|j durX| |nd} |so| durd|| fn|f} | |ddSt || |j |jdS)Nz You have to specify pixel_valuesrrK)rTrUrrrrrr)r& pooler_outputr'r()r3rruse_return_dictr get_head_maskrrVr rNr r rr'r() rGrJrrrrembedding_outputencoder_outputssequence_output pooled_output head_outputsr0r0r1rZgs:  zYolosModel.forward)T)NNNNN)r)r*r+rrr7r>rrrrrrYOLOS_INPUTS_DOCSTRINGr_CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPErr-r rrrZr\r0r0rHr1rGs<    rcs*eZdZdeffdd ZddZZS)r r3cs*tt|j|j|_t|_dSr_)r6r7r rr:rTanh activationrarHr0r1r7s zYolosPooler.__init__cCs(|dddf}||}||}|S)Nr)rr%)rGr'first_token_tensorrr0r0r1rZs  zYolosPooler.forward)r)r*r+rr7rZr\r0r0rHr1r sr zy YOLOS Model (consisting of a ViT encoder) with object detection heads on top, for tasks such as COCO detection. cseZdZdeffdd ZejjddZe e e e e d    ddejd eeed eed eed eed eee ff ddZZS)YolosForObjectDetectionr3csXt|t|dd|_t|j|j|jddd|_t|j|jddd|_| dS)NF)rrr ) input_dim hidden_dim output_dim num_layers) r6r7rrYolosMLPPredictionHeadr: num_labelsclass_labels_classifierbbox_predictorr rarHr0r1r7s   z YolosForObjectDetection.__init__cCs$ddt|dd|ddDS)NcSsg|] \}}||dqS))r#r$r0)rabr0r0r1rsz9YolosForObjectDetection._set_aux_loss..rK)zip)rG outputs_class outputs_coordr0r0r1 _set_aux_losss$z%YolosForObjectDetection._set_aux_loss)rrNrJlabelsrrrr4c s|dur|n|jj}|j||||d}|d}|dd|jj dddf}||}||} d\} } |durt|jj|jj |jj d} gd} t | |jj |jj | d}||ji}||d<| |d <|jjr|rr|jn|d }||}||}|||} | |d <|||d |jjd |jjd<|jjri}t|jjd D]|fddDq|tfddD} |s| dur|| f| |}n|| f|}| dur| 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_cost)r7boxes cardinality)matcher num_classeseos_coeflossesr#r$r,r%r)loss_ce loss_bbox loss_giouc i|] \}}|d|qSrXr0rkrrr0r1 = z3YolosForObjectDetection.forward..c3s(|]}|vr||VqdSr_r0)rrG)r" weight_dictr0r1r?s&z2YolosForObjectDetection.forward..)r!r"r#r$r%r&r'r()r3rrr<r/r0sigmoidYolosHungarianMatcherr8r9r: YolosLossr.eos_coefficienttodeviceauxiliary_lossintermediate_hidden_statesr6bbox_loss_coefficientgiou_loss_coefficientrdecoder_layersupdatersumkeysr r&r'r()rGrJr7rrrrrr#r$r!r%r=r@ criterion outputs_lossrr4r5aux_weight_dictrr0)rr"rKr1rZsr7          zYolosForObjectDetection.forward)NNNN)r)r*r+rr7r-jitunusedr6rr rr r"r.rrrrrrrZr\r0r0rHr1r's.    r'cCsX|}|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). rrdrK)rLrorX)inputstargets num_boxes numerator denominatorr!r0r0r1 dice_lossUs   rd?rdalphagammac Cs||}tjj||dd}||d|d|}|d||}|dkr5||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 none reductionrr)rLr rm binary_cross_entropy_with_logitsrrX) r_r`rarfrgprobce_lossp_tr!alpha_tr0r0r1sigmoid_focal_lossisrpcsheZdZdZfddZddZeddZdd Z d d Z d d Z ddZ ddZ ddZZS)rNa 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)NrrK empty_weight) r6r7r=r>r?r@r-onesregister_buffer)rGr=r>r?r@rqrHr0r1r7s  zYolosLoss.__init__c Csd|vrtd|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 outputscS g|] \}\}}|d|qS class_labelsr0)rtrXJr0r0r1rrJz)YolosLoss.loss_labels..NrddtyperQrrA)KeyError_get_source_permutation_idxr-rQr3fullrNr>int64rQr rm cross_entropyrirq) rGrr`indicesra source_logitsidxtarget_classes_otarget_classesrAr@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#cSg|]}t|dqSrurrr0r0r1rz.YolosLoss.loss_cardinality..)rQrKrcardinality_error) rQr- as_tensorargmaxrNrXr rml1_lossfloat) rGrr`rrar#rQtarget_lengths card_predcard_errr@r0r0r1loss_cardinalityszYolosLoss.loss_cardinalityc Csd|vrtd||}|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 outputscSrtr;r0)rrwrXrr0r0r1rrJz(YolosLoss.loss_boxes..rrLrhrirBrrC) r{r|r-rQr3r rmrrXdiaggeneralized_box_iour) rGrr`rrar source_boxes target_boxesrBr@rCr0r0r1 loss_boxess  zYolosLoss.loss_boxesc Csd|vrtd||}||}|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]. pred_masksz#No predicted masks found in outputscSg|]}|dqS)masksr0rrwr0r0r1rrz(YolosLoss.loss_masks..NrbilinearFrfrr) loss_mask loss_dice)r{r|_get_target_permutation_idxnested_tensor_from_tensor_list decomposerPr rmrnrNrorlrprd) rGrr`rra source_idx target_idx source_masksr target_masksvalidr@r0r0r1 loss_maskss(       zYolosLoss.loss_maskscC4tddt|D}tdd|D}||fS)NcSs g|] \}\}}t||qSr0r- full_like)rrsourcerXr0r0r1r rJz9YolosLoss._get_source_permutation_idx..cSsg|]\}}|qSr0r0)rrrXr0r0r1rrr-rQr)rGr batch_idxrr0r0r1r| z%YolosLoss._get_source_permutation_idxcCr)NcSs g|] \}\}}t||qSr0r)rrrXtargetr0r0r1rrJz9YolosLoss._get_target_permutation_idx..cSsg|]\}}|qSr0r0)rrXrr0r0r1rrr)rGrrrr0r0r1rrz%YolosLoss._get_target_permutation_idxcCs@|j|j|j|jd}||vrtd|d||||||S)N)r7r<r;rzLoss z not supported)rrrrr)rGr!rr`rraloss_mapr0r0r1get_lossszYolosLoss.get_lossc sdd|D}|||}tdd|D}tj|gtjtt|j d}tj |dd }i}|j D]}| ||||||q7d|vrt|dD].\}|||}|j D] }|d krdq]||||||} fd d| D} | | q]qP|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%r0rFr0r0r1rI-sz%YolosLoss.forward..css|] }t|dVqdS)rvNrrr0r0r1r3rz$YolosLoss.forward..ryrminr%rcrDrEr0rFrHr0r1rIIrJ)rr=rXr-rrnextitervaluesrQclampitemr@rWrr) rGrr`outputs_without_auxrrar@r!r%l_dictr0rHr1rZ"s& "    zYolosLoss.forward)r)r*r+r,r7rr-no_gradrrrr|rrrZr\r0r0rHr1rNs   ! rNcs(eZdZdZfddZddZZS)r-a 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)rnrGr0r0r1r]sz2YolosMLPPredictionHead.__init__..)r6r7r+r rr3layers)rGr(r)r*r+hrHr0r1r7Ys ,zYolosMLPPredictionHead.__init__cCs>t|jD]\}}||jdkrtj||n||}q|Sr5)rrr+r rmrelu)rGrrrr0r0r1rZ_s(zYolosMLPPredictionHead.forward)r)r*r+r,r7rZr\r0r0rHr1r-Ps r-cs@eZdZdZd dededeffdd Zedd ZZ S) rMa 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. rr8r9r:csTtt|dg||_||_||_|dkr$|dkr&|dkr(tddSdSdS)Nscipyrz#All costs of the Matcher can't be 0)r6r7rr8r9r:r)rGr8r9r:rHr0r1r7ws   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#NrdrrrKr$cSrrur0rr0r0r1rrz1YolosHungarianMatcher.forward..cSrrr0rr0r0r1rr)pcSrrrrr0r0r1rrcSsg|] \}}t||qSr0r)rrcr0r0r1rscSs0g|]\}}tj|tjdtj|tjdfqS))rz)r-rr~)rrjr0r0r1rs0)rNrorr-rQcdistrrr9r8r:rlcpursplit)rGrr`rR num_queriesout_probout_bbox target_ids target_bboxr8r9r: cost_matrixsizesrr0r0r1rZszYolosHungarianMatcher.forward)rrr) r)r*r+r,rr7r-rrZr\r0r0rHr1rMfs  rMrwr4cCsD|r|jtjtjfvr|S|S|jtjtjfvr|S|Sr_) is_floating_pointrzr-float32float64rint32r~r)rwr0r0r1_upcastsrr;cCsHt|}|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. Nrdrr r)rrr0r0r1box_areas @rc Cst|}t|}t|dddddf|ddddf}t|dddddf|ddddf}||jdd}|dddddf|dddddf}|dddf||}||} | |fS)Nrdrrr)rr-maxrr) boxes1boxes2area1area2left_top right_bottom width_heightinterriour0r0r1box_ious..,rcCs*|ddddf|ddddfkstd||ddddf|ddddfks:td|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) Nrdz.ryrrdFz(Only 3-dimensional tensors are supported)ndimrrrzrQr-r9rrrr3rNcopy_rr)rmax_size batch_shaperRrSrTrUrzrQtensorrrpad_imgmr0r0r1rs   2" r)rerd)Sr,collections.abcrr dataclassesrtypingrrrrrrr-torch.utils.checkpointr r activationsr modeling_outputsr rmodeling_utilsr pytorch_utilsrrutilsrrrrrrrrrconfiguration_yolosrscipy.optimizertransformers.image_transformsr get_loggerr)loggerr"r!r##YOLOS_PRETRAINED_MODEL_ARCHIVE_LISTr Moduler2rEr~r>rrrrrrrrYOLOS_START_DOCSTRINGr rr r'rdrrprNr-rMrrrrrobjectrrr0r0r0r1s    ,    ,+!"=(*N W HO