|h;XddlmZmZmZmZmZddlZddlmZddl mcm Z ddl m Z ddlmZddlmZmZGddej&Z ddeeefd ed ed ej.d ed edededeeej.eej.eej.eeeefffdZy))AnyDictListOptionalTupleN)linear_sum_assignment)bbox_iou) xywh2xyxy xyxy2xywhcPeZdZdZ ddeeeefdedede dedef fd Z dd e jd e jd e jd e jd e e dee jdee e jde ee je jffdZxZS)HungarianMatchera A module implementing the HungarianMatcher for optimal assignment between predictions and ground truth. HungarianMatcher performs optimal bipartite assignment over predicted and ground truth bounding boxes using a cost function that considers classification scores, bounding box coordinates, and optionally mask predictions. This is used in end-to-end object detection models like DETR. Attributes: cost_gain (Dict[str, float]): Dictionary of cost coefficients for 'class', 'bbox', 'giou', 'mask', and 'dice' components. use_fl (bool): Whether to use Focal Loss for classification cost calculation. with_mask (bool): Whether the model makes mask predictions. num_sample_points (int): Number of sample points used in mask cost calculation. alpha (float): Alpha factor in Focal Loss calculation. gamma (float): Gamma factor in Focal Loss calculation. Methods: forward: Compute optimal assignment between predictions and ground truths for a batch. _cost_mask: Compute mask cost and dice cost if masks are predicted. Examples: Initialize a HungarianMatcher with custom cost gains >>> matcher = HungarianMatcher(cost_gain={"class": 2, "bbox": 5, "giou": 2}) Perform matching between predictions and ground truth >>> pred_boxes = torch.rand(2, 100, 4) # batch_size=2, num_queries=100 >>> pred_scores = torch.rand(2, 100, 80) # 80 classes >>> gt_boxes = torch.rand(10, 4) # 10 ground truth boxes >>> gt_classes = torch.randint(0, 80, (10,)) >>> gt_groups = [5, 5] # 5 GT boxes per image >>> indices = matcher(pred_boxes, pred_scores, gt_boxes, gt_classes, gt_groups) cost_gainuse_fl with_masknum_sample_pointsalphagammact||dddddd}||_||_||_||_||_||_y)a Initialize HungarianMatcher for optimal assignment of predicted and ground truth bounding boxes. Args: cost_gain (Dict[str, float], optional): Dictionary of cost coefficients for different matching cost components. Should contain keys 'class', 'bbox', 'giou', 'mask', and 'dice'. use_fl (bool): Whether to use Focal Loss for classification cost calculation. with_mask (bool): Whether the model makes mask predictions. num_sample_points (int): Number of sample points used in mask cost calculation. alpha (float): Alpha factor in Focal Loss calculation. gamma (float): Gamma factor in Focal Loss calculation. N)classbboxgioumaskdice)super__init__rrrrrr)selfrrrrrr __class__s [/var/www/html/test/engine/venv/lib/python3.12/site-packages/ultralytics/models/utils/ops.pyrzHungarianMatcher.__init__0sQ*   "#Q1aPI" "!2  pred_bboxes pred_scores gt_bboxesgt_cls gt_groupsmasksgt_maskreturnc  |j\}} } t|dk(rat|D cgc]L} tjgtj tjgtj fNc} S|j jd| }|jrtj|ntj|d}|j jdd}|dd|f}|jrqd|jz ||jzzd|z dzj z} |jd|z |jzz|dzj z} | | z }n| }|jd|jdz j!jd}d t#|jd|jdd d j%dz }|j&d |z|j&d |zz|j&d|zz}|j(r||j+||||z }d||j-|j/z<|j|| dj1}t3|j5|dDcgc]\}}t7||}}}tj8dg|ddj;d}t3|Dcgc]X\}\}}tj|tj tj|tj ||zfZc}}}Scc} wcc}}wcc}}}w)a Compute optimal assignment between predictions and ground truth using Hungarian algorithm. This method calculates matching costs based on classification scores, bounding box coordinates, and optionally mask predictions, then finds the optimal bipartite assignment between predictions and ground truth. Args: pred_bboxes (torch.Tensor): Predicted bounding boxes with shape (batch_size, num_queries, 4). pred_scores (torch.Tensor): Predicted classification scores with shape (batch_size, num_queries, num_classes). gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (num_gts, 4). gt_cls (torch.Tensor): Ground truth class labels with shape (num_gts,). gt_groups (List[int]): Number of ground truth boxes for each image in the batch. masks (torch.Tensor, optional): Predicted masks with shape (batch_size, num_queries, height, width). gt_mask (List[torch.Tensor], optional): Ground truth masks, each with shape (num_masks, Height, Width). Returns: (List[Tuple[torch.Tensor, torch.Tensor]]): A list of size batch_size, each element is a tuple (index_i, index_j), where index_i is the tensor of indices of the selected predictions (in order) and index_j is the tensor of indices of the corresponding selected ground truth targets (in order). For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes). rdtypedimNrg:0yE>?T)xywhGIoUrrr)shapesumrangetorchtensorlongdetachviewrFsigmoidsoftmaxrrlog unsqueezeabsr squeezerr _cost_maskisnanisinfcpu enumeratesplitr as_tensorcumsum_)rr#r$r%r&r'r(r)bsnqnc_neg_cost_classpos_cost_class cost_class cost_bbox cost_giouCicindiceskjs r!forwardzHungarianMatcher.forwardOs'@!&& B y>Q fklnfopabU\\"EJJ7bPUPZPZ9[\p p"((*//B7 04 aii ,;\^A_ !((*//A6 "!V), ;;$**ndjj1HIqS^aeOeNjNjNlMlmN!ZZA O +JKQ\_cQcPhPhPjOjkN'.8J%J!**1- 0C0CA0FFKKMQQRTU (;#8#8#;Y=P=PQR=SZ^eijrrsuvv  NN7 #j 0nnV$y0 1nnV$y0 1  >> Yw? ?A$'!'')aggi  FF2r2  " " $;DQWWYXZE[;\]41a(1.]]OOQ$83B$89AA!D 'w/  6Aq\\!5:: . Qejj0QT]^_T`0` a  OqJ^ sAM4>M9AM?)NTFi1g?@)NN)__name__ __module__ __qualname____doc__rrstrfloatboolintrr9Tensorrrr\ __classcell__)r s@r!r r sF15!&De,-     L)-04M \\M \\M << M  M 9 M  %M $u||,-M  eELL%,,./ 0M r"r batch num_classes num_queries class_embednum_dncls_noise_ratiobox_noise_scaletrainingr*c  |r|dks|y|d}t|} t|} | dk(ry|| z} | dk(rdn| } t|} |d} |d}|d}| jd | z}|jd | zd}|jd | zj d }t j | | zt j|j | | zz}|dkDr|t j|j|d zk}t j|jd }t j|d||j|j }|||<|dkDrt|}|d d dfd zjdd |z}t j|dd dzdz }t j |}||xxdz cc<||z}|||zz }|j#ddt%|}t j&|d}t)| d z| z}||}t j*| ||jd | j}t j*| |d|j}t j,|Dcgc]0}t j.t1|t j2c}}t j2t1| D cgc] } || | zz c} d}!t j,t1d | zD cgc] } || | zz c} }||||f<||||f<||z}"t j*|"|"gt j4}#d|#|dd|f<t1| D]} | dk(r#d|#| d z| z| d z| dzz| d z| dzz|f<| | dz k(r!d|#| d z| z| d z| dzzd| | zd zf<Td|#| d z| z| d z| dzz| d z| dzz|f<d|#| d z| z| d z| dzzd| d z| zf<|!j7j9t;|dD$cgc]}$|$j=d c}$| ||gd}%|j?|j|j?|j|#j?|j|%fScc}wcc} wcc} wcc}$w)a Generate contrastive denoising training group with positive and negative samples from ground truths. This function creates denoising queries for contrastive denoising training by adding noise to ground truth bounding boxes and class labels. It generates both positive and negative samples to improve model robustness. Args: batch (Dict[str, Any]): Batch dictionary containing 'gt_cls' (torch.Tensor with shape (num_gts,)), 'gt_bboxes' (torch.Tensor with shape (num_gts, 4)), and 'gt_groups' (List[int]) indicating number of ground truths per image. num_classes (int): Total number of object classes. num_queries (int): Number of object queries. class_embed (torch.Tensor): Class embedding weights to map labels to embedding space. num_dn (int): Number of denoising queries to generate. cls_noise_ratio (float): Noise ratio for class labels. box_noise_scale (float): Noise scale for bounding box coordinates. training (bool): Whether model is in training mode. Returns: padding_cls (torch.Tensor | None): Modified class embeddings for denoising with shape (bs, num_dn, embed_dim). padding_bbox (torch.Tensor | None): Modified bounding boxes for denoising with shape (bs, num_dn, 4). attn_mask (torch.Tensor | None): Attention mask for denoising with shape (tgt_size, tgt_size). dn_meta (Dict[str, Any] | None): Meta information dictionary containing denoising parameters. Examples: Generate denoising group for training >>> batch = { ... "cls": torch.tensor([0, 1, 2]), ... "bboxes": torch.rand(3, 4), ... "batch_idx": torch.tensor([0, 0, 1]), ... "gt_groups": [2, 1], ... } >>> class_embed = torch.rand(80, 256) # 80 classes, 256 embedding dim >>> cdn_outputs = get_cdn_group(batch, 80, 100, class_embed, training=True) rN)NNNNr'rclsbboxes batch_idxrr.)r-device?.r]r2r5)minmaxgư>)eps)rtr1r,r/T) dn_pos_idx dn_num_group dn_num_split) r7rwlenrepeatr=r9aranger;rtrandr6nonzerorD randint_liker-r rand_likeclip_r logitrezeroscatr:r8stackrdrHrJlistreshapeto)&rhrirjrkrlrmrnror' total_nummax_nums num_grouprMr&gt_bboxb_idxdn_clsdn_bboxdn_b_idxneg_idxridx new_label known_bboxdiff rand_sign rand_part dn_cls_embed padding_cls padding_bboxnum map_indicesrWpos_idxtgt_size attn_maskpdn_metas& r! get_cdn_grouprsZ 1 %k"III9~H1}%("I!^I YB 5\FHoG + E]]1y= )FnnQ]A.G||A M*//3Hll9y0 7>>Z]fir]rrGzz&,,'?S+@AmmD!))"-&&sA{&,,W]WdWde s w' QR 3&..q!4F&&w15;cA OOG, 'c!Y i$&& Sc*J'++g40 A ) *Fv&L++b&,*<*scipy.optimizerultralytics.utils.metricsr ultralytics.utils.opsr r Moduler rbrerfrcrdrr"r!rs43 0.6N ryyN h  @ S>@@@ @  @  @@@ 8ELL !8ELL#98ELL;QS[\`adfiai\jSk kl@r"