o Zh >@sddlZddlmZddlmmZddlmZddlm Z ddl m Z m Z er0ddl mZejjdd Zejj dd d Zd ed ejdedejfddZdejd eded ejfddZdddZGddde Z       dddZdS)N)is_vision_available)box_iou)RTDetrHungarianMatcher RTDetrLoss)center_to_corners_formatcCsddt||DS)NcSsg|] \}}||dqS))logits pred_boxes).0abr r L/var/www/auris/lib/python3.10/site-packages/transformers/loss/loss_d_fine.py $z!_set_aux_loss..zip) outputs_class outputs_coordr r r _set_aux_losssrcsfddt||||DS)Nc s&g|]\}}}}||||dqS))r r pred_corners ref_pointsteacher_cornersteacher_logitsr )r r rcdrrr rr.s z"_set_aux_loss2..r)rrZoutputs_cornersZ outputs_refrrr rr_set_aux_loss2's  r max_num_binsup reg_scalereturncst|dt|}t|dt|d}|dd|dfddt|ddddD}fddtd|dD}| g|t|ddg||g}d d|D}t|d}|S) uK Generates the non-uniform Weighting Function W(n) for bounding box regression. Args: max_num_bins (int): Max number of the discrete bins. up (Tensor): Controls upper bounds of the sequence, where maximum offset is ±up * H / W. reg_scale (float): Controls the curvature of the Weighting Function. Larger values result in flatter weights near the central axis W(max_num_bins/2)=0 and steeper weights at both ends. Returns: Tensor: Sequence of Weighting Function. rrrcsg|] }| dqSrr r istepr rrLrz&weighting_function..csg|]}|dqSr#r r$r&r rrMsNcSs&g|]}|dkr |n|dqS)r)dim unsqueeze)r vr r rrOs&)absrangetorch zeros_likecat)rr r!Z upper_bound1Z upper_bound2 left_values right_valuesvaluesr r&rweighting_function;s"* r4gtcCs"|d}t|||}|d|d}|dk}tj|ddd}|}t|} t|} |dk||k@} || } || } || d}t|| | }t||| }|||| | <d| | | | <|dk}d| |<d| |<d||<||k}d| |<d| |<|d||<|| | fS)a  Decodes bounding box ground truth (GT) values into distribution-based GT representations. This function maps continuous GT values into discrete distribution bins, which can be used for regression tasks in object detection models. It calculates the indices of the closest bins to each GT value and assigns interpolation weights to these bins based on their proximity to the GT value. Args: gt (Tensor): Ground truth bounding box values, shape (N, ). max_num_bins (int): Maximum number of discrete bins for the distribution. reg_scale (float): Controls the curvature of the Weighting Function. up (Tensor): Controls the upper bounds of the Weighting Function. Returns: Tuple[Tensor, Tensor, Tensor]: - indices (Tensor): Index of the left bin closest to each GT value, shape (N, ). - weight_right (Tensor): Weight assigned to the right bin, shape (N, ). - weight_left (Tensor): Weight assigned to the left bin, shape (N, ). r(rrr)g?g皙?) reshaper4r*r.sumfloatr/longr,)r5rr!r Zfunction_valuesZdiffsmaskZclosest_left_indicesindices weight_right weight_leftZvalid_idx_maskZ valid_indicesr1r2Z left_diffsZ right_diffsZinvalid_idx_mask_negZinvalid_idx_mask_posr r r translate_gtTs2        r@r7c CsFt|}|dddf|dddf|d|dd|}|dddf|dddf|d|dd|}|dddf|dddf|d|dd|}|ddd f|dddf|d|dd|} t|||| gd } t| |||\} } } |dur| jd||d } | d | | fS) a Converts bounding box coordinates to distances from a reference point. Args: points (Tensor): (n, 4) [x, y, w, h], where (x, y) is the center. bbox (Tensor): (n, 4) bounding boxes in "xyxy" format. max_num_bins (float): Maximum bin value. reg_scale (float): Controlling curvarture of W(n). up (Tensor): Controlling upper bounds of W(n). eps (float): Small value to ensure target < max_num_bins. Returns: Tensor: Decoded distances. Nr).rgؗҜ<?r).rrBr(minmax)r,r.stackr@clampr8detach) ZpointsZbboxrr!r ZepslefttoprightbottomZ four_lensr>r?r r r bbox2distances8888rMcs>eZdZdZfddZ d ddZdd d Zd d ZZS) DFineLossa This class computes the losses for D-FINE. 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). Args: matcher (`DetrHungarianMatcher`): Module able to compute a matching between targets and proposals. weight_dict (`Dict`): Dictionary relating each loss with its weights. These losses are configured in DFineConf as `weight_loss_vfl`, `weight_loss_bbox`, `weight_loss_giou`, `weight_loss_fgl`, `weight_loss_ddf` losses (`List[str]`): List of all the losses to be applied. See `get_loss` for a list of all available losses. alpha (`float`): Parameter alpha used to compute the focal loss. gamma (`float`): Parameter gamma used to compute the focal loss. eos_coef (`float`): Relative classification weight applied to the no-object category. num_classes (`int`): Number of object categories, omitting the special no-object category. csjt|t||_|j|_|j|j|j|j|j d|_ gd|_ |j |_ t jt|jgdd|_dS)N)Zloss_vflZ loss_bboxZ loss_giouloss_fglloss_ddf)vflboxeslocalF)Z requires_grad)super__init__rZmatcherrZweight_loss_vflZweight_loss_bboxZweight_loss_giouZweight_loss_fglZweight_loss_ddfZ weight_dictlossesr!nn Parameterr.Ztensorr )selfconfig __class__r rrUs   zDFineLoss.__init__Nr9c Cs|}|d} tj||dd|dtj|| dd|d} |dur.|}| |} |dur:| |} | S|dkrD| } | S|dkrL| } | S)Nrnone reductionr(meanr9)r;FZ cross_entropyr8r:r9r`) rYpredlabelr>r?weightr_ avg_factorZdis_leftZ dis_rightlossr r r unimodal_distribution_focal_losss& z*DFineLoss.unimodal_distribution_focal_losscCsi}d|vrZ||}tjddt||Ddd}|d|d|jd} |d|} tt| t ||j|j |j |_ Wd n1sMwY|j \} } } t tt |d |t |d}|dddd d}|j| | | | ||d |d <|dd|jd} |dd|jd} t| | r| d|d<|S|djddd}tj|tjd}d||<|dddd d}||||j||<|dddd d}||dtjddtj| |ddtj| |ddd}d|d j d}||d||d|_!|_"|#r7||$nd}|#rF||$nd}||j!||j"|j!|j"|d<|S)zaCompute Fine-Grained Localization (FGL) Loss and Decoupled Distillation Focal (DDF) Loss.rcSs g|] \}\}}|d|qS)rRr )r t_r%r r rrs z(DFineLoss.loss_local..rr6r(rrNr )rerOrrPr)dtypeTrr]r^rA)%Z_get_source_permutation_idxr.r0rr8rrHZno_gradrMrr!r Z fgl_targetsZdiagrr*repeatrgequalr9ZsigmoidrEr/boolZ reshape_astorlrWZ KLDivLossraZ log_softmaxZsoftmaxshapenum_posZnum_neganyr`)rYoutputstargetsr= num_boxesTrVidxZ target_boxesrrZtarget_cornersr>r?ZiousZweight_targetsZweight_targets_localr<Zloss_match_localZ batch_scaleZloss_match_local1Zloss_match_local2r r r loss_localsz            zDFineLoss.loss_localcCsD|j|j|j|j|jd}||vrtd|d||||||S)N)Z cardinalityrSrRZfocalrQzLoss z not supported)Zloss_cardinalityryZ loss_boxesZloss_labels_focalZloss_labels_vfl ValueError)rYrfrtrur=rvZloss_mapr r rget_loss9szDFineLoss.get_loss)Nr9N)rh) __name__ __module__ __qualname____doc__rUrgryr{ __classcell__r r r[rrNs  KrNc  Kst|} | |i}||d<|jddd|d<d}|jr| durtj|jddd| ddd\}}tj|| ddd\}}tj| | ddd\}}tj| | ddd\}}t|dddd fdd|dddd fdd|dddd fdd|dddd fdd|ddd f|ddd f}||d <|d t |g|jdddgt|dd|dd|dd|dd|ddd f|ddd f}||d <| |d <| ||}t | }|||fS) Nr rrrCr Z dn_num_splitrr6r(auxiliary_outputsdn_auxiliary_outputsdenoising_meta_values) rNrprGZauxiliary_lossr.splitrZ transposeextendrr9r3)r labelsZdevicer rZrrZenc_topk_logitsZenc_topk_bboxesrZpredicted_cornersZinitial_reference_pointskwargs criterionZ outputs_lossrZ dn_out_coordZ dn_out_classZdn_out_cornersZ out_cornersZ dn_out_refsZout_refsrZ loss_dictrfr r rDFineForObjectDetectionLossFsN          r)NN)r7)NNNNNNN)r.Ztorch.nnrWZtorch.nn.functionalZ functionalrautilsrZloss_for_object_detectionrZ loss_rt_detrrrZtransformers.image_transformsrZjitZunusedrrintZTensorr4r@rMrNrr r r rs4       ?