FPhdrddlZddlmZddlmcmZddlmZddlm Z m Z m Z ddl m Z mZmZddlmZddlmZGdd ej(ZGd d ej(ZGd d ej(ZGddej(ZGddZGddeZGddeZGddZy)N) OKS_SIGMA) crop_mask xywh2xyxy xyxy2xywh)TaskAlignedAssigner dist2bbox make_anchors)bbox_iou) bbox2distc4eZdZdZfdZeddZxZS) VarifocalLosszO Varifocal loss by Zhang et al. https://arxiv.org/abs/2008.13367. c"t|y)z#Initialize the VarifocalLoss class.Nsuper__init__self __class__s aC:\Users\daisl\Desktop\realtime-object-detection\venv\Lib\site-packages\ultralytics/utils/loss.pyrzVarifocalLoss.__init__ c||jj|zd|z z||zz}tjjj d5t j|j|jd|zjdj}ddd|S#1swYSxYw)zComputes varfocal loss.r F)enablednone reductionN) sigmoidpowtorchcudaampautocastF binary_cross_entropy_with_logitsfloatmeansum) pred_scoregt_scorelabelalphagammaweightlosss rforwardzVarifocalLoss.forwards++-11%88AIFTYIYY ZZ^^ $ $U $ 366z7G7G7I8>>K[gmn DGCCE 4 4 s AB88C)g?@__name__ __module__ __qualname____doc__r staticmethodr0 __classcell__rs@rrrs! rrc4eZdZdZfdZeddZxZS) FocalLosszjWraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5).c"t|y)z3Initializer for FocalLoss class with no parameters.Nrrs rrzFocalLoss.__init__'rrctj||d}|j}||zd|z d|z zz}d|z |z}||z}|dkDr||zd|z d|z zz}||z}|jdj S)zRCalculates and updates confusion matrix for object detection/classification tasks.rrr ?r)r$r%rr'r() predr+r-r,r/ pred_probp_tmodulating_factor alpha_factors rr0zFocalLoss.forward+s11$P LLN i1u9Y"?? 3Y50 !! 19 5=AI!e)+DDL L Dyy|!!r)g?g?r2r9s@rr;r;$st""rr;c:eZdZdZdfd ZdZedZxZS)BboxLossz>Criterion class for computing training losses during training.c>t|||_||_y)zLInitialize the BboxLoss module with regularization maximum and DFL settings.N)rrreg_maxuse_dfl)rrGrHrs rrzBboxLoss.__init__@s   rc|jd|jd}t||||dd} d| z |zj|z } |jrft |||j } |j ||jd|j dz| ||z} | j|z } | | fStjdj|j} | | fS)z IoU loss.FT)xywhCIoUr>r g) r( unsqueezer rHr rG_df_lossviewr tensortodevice) r pred_dist pred_bboxes anchor_points target_bboxes target_scorestarget_scores_sumfg_maskr.iouloss_iou target_ltrbloss_dfls rr0zBboxLoss.forwardFs""2&w/99"={7+]7-C%VZ[3Y&(--/2CC <<#M=$,,OK}}Yw%7%<%rJ)rcr r(rmrOexpr')r pred_kptsgt_kptskpt_maskareadkpt_loss_factores rr0zKeypointLoss.forwardjs v  0Q 6)F:KgV\o:]bc9c c"..+uyyQA/NQU/UV T[[Q& &$+ 6 :$$R+EIIqbM0AX/MNTTVVr)returnN)r3r4r5r6rr0r8r9s@rrjrjbs8 Wrrjc(eZdZdZdZdZdZdZy)v8DetectionLossrkct|jj}|j}|jd}t j d|_||_|j|_ |j|_ |j|_ |j|_ ||_|jdkD|_ td|jdd|_t!|jdz |j j#||_t'j(|jt&j*| |_y ) zdInitializes v8DetectionLoss with the model, defining model-related properties and BCE loss function.rJrrr ?g@)topk num_classesr,beta)rH)dtyperRN)next parametersrRargsmodelnnBCEWithLogitsLossbcehypstridencnorGrHrassignerrErQ bbox_lossr aranger&proj)rrrRhms rrzv8DetectionLoss.__init__vse&&()00 JJ KKO''&9hh $$$$yy  yy1} +PSZ]^ !!))a-FII&QLL%++fM rc|jddk(r%tj|dd|j}|S|dddf}|j d\}}|j tj }tj||jd|j}t|D]*}||k(} | j} | s|| ddf||d| f<,t|d ddfj||d ddf<|S) zXPreprocesses the target counts and matches with the input batch size to output a tensor.rrRNT return_counts)rr .) rcr zerosrRuniquerQint32maxranger(rmul_) rtargets batch_size scale_tensorouti_countsjmatchesns r preprocesszv8DetectionLoss.preprocesss == q ++j!Qt{{CC 1 At4IAvYYU[[Y1F++j&**,$++NC:&q&KKM!(!"!5C2A2J ' &c#qs(m&8&8&FGCQqSM rc|jrh|j\}}}|j||d|dzjdj |j j |j}t||dS)zUDecode predicted object bounding box coordinates from anchor points and distribution.F)rK) rHrcrOsoftmaxmatmulrtyperr)rrUrSbacs r bbox_decodezv8DetectionLoss.bbox_decodesk <<ooGAq!!q!QQ7??BII$))..YbYhYhJijIM>>rc <tjd|j}t|tr|dn|}tj |Dcgc]/}|j |djd|jd1c}dj|jdz|jfd\}}|jdddj}|jdddj}|j}|jd} tj|djdd|j| |j dz} t#||j d \} } tj |d j dd|d j dd|d fd} |j%| j'|j| | gd} | jdd\}}|j)ddj+d}|j-| |}|j/|j1j3|j1| zj5|j| | z|||\}}}}}t7|j)d}|j9||j'|j)|z |d<|j)r%|| z}|j;||| ||||\|d<|d<|dxx|j<j>zcc<|dxx|j<j@zcc<|dxx|j<jBzcc<|j)| z|j1fScc}w)zLCalculate the sum of the loss for box, cls and dfl multiplied by batch size.rrr rrJrprNrRrr batch_idxclsbboxesr rr rrr rTr_)"r rrR isinstancetuplecatrOrcrsplitrGrpermute contiguousrrPrr rrQr(gt_rrdetachrrrrrrboxrdfl)rpredsbatchr/featsxi pred_distri pred_scoresrrimgszrU stride_tensorr gt_labels gt_bboxesmask_gtrTrrVrWrYrXs r__call__zv8DetectionLoss.__call__s<{{1T[[1&ue4a%#(99`e-f`eZ\bggeAhnnQ6GRT.U`e-fhi#j#p#p \\A tww '$, ["))!Q2==? !))!Q2==? !! &&q)  U1X^^AB/ 5QTXT_T_`aTbb'3E4;;'L$ }))U;/44R;U5\=N=NrST=UW\]eWfgijk//'**T[["9:TYZfTg/h&}}VQ7 9--4-044Q7&&}kB 6:mm    ( ( *[-?-?-AM-Q,W,WXaXgXg,h M )9i7J3=-! 1 1 3Q7((; (8(8(?@DDFIZZQ ;;= ] *M#~~k; Wdfs.? J DGT!W Q488<< Q488<< Q488<<xxzJ& 55Q.gs 4NN)r3r4r5r6rrrrrrr~r~ss8N("?,6rr~ceZdZdZfdZdZedejdejdejdejdejd ejf d Z d ejd ejd ejdejdejdejdejdejde d ejfdZ xZ S)v8SegmentationLossrkcZt|||jj|_y)zSInitializes the v8SegmentationLoss class, taking a de-paralleled model as argument.N)rrr overlap_maskoverlap)rrrs rrzv8SegmentationLoss.__init__s! zz.. rc  tjd|j}t|dk(r|n|d\}}}|j\}}} } tj |D cgc]/} | j |djd|jd1c} dj|jdz|jfd\} } | jdddj} | jdddj} |jdddj}| j}tj|djdd|j| |jdz}t!||jd \}} |d j dd}tj ||d j dd|d fd}|j#|j%|j||gd}|jdd\}}|j'ddj)d}|j/|| }|j1| j3j5|j3|zj7|j||z|||\}}}}}t9|j'd}|j;| |j%|j'|z |d<|j'r|j=| ||||z |||\|d<|d<|dj%|jj?}tA|jdd| | fk7r tCjD|d| | fdd}|jG|||||||||jH |d<n2|dxx|dzj'|dzj'zz cc<|dxx|jJjLzcc<|dxx|jJjLzcc<|dxx|jJjNzcc<|dxx|jJjPzcc<|j'|z|j3fScc} w#t*$r}t-d|d}~wwxYw)z1Calculate and return the loss for the YOLO model.rrrr rrJrpNrrrrrrrrTr_uERROR ❌ segment dataset incorrectly formatted or not a segment dataset. This error can occur when incorrectly training a 'segment' model on a 'detect' dataset, i.e. 'yolo train model=yolov8n-seg.pt data=coco128.yaml'. Verify your dataset is a correctly formatted 'segment' dataset using 'data=coco128-seg.yaml' as an example. See https://docs.ultralytics.com/tasks/segment/ for help.masksnearest)mode))r rrRlenrcrrOrrrGrrrrrPrr rrQr(r RuntimeError TypeErrorrrrrrrrrr&rr$ interpolatecalculate_segmentation_lossrrrrr)rrrr/r pred_masksprotorrmask_hmask_wrrrrrrUrrrrrrr{rTrVrWrY target_gt_idxrXrs rrzv8SegmentationLoss.__call__se{{1T[[1,/J!O5q z5(- % Avv#(99`e-f`eZ\bggeAhnnQ6GRT.U`e-fhi#j#p#p \\A tww '$, ["))!Q2==? !))!Q2==? ''1a0;;= !! U1X^^AB/ 5QTXT_T_`aTbb'3E4;;'L$ } pk*//A6IiiE%L,=,=b!,DeHo VXYZGoogjj&=zX]^jXkolG#*==#; IymmAtm488;G&&}kB BF--    ( ( *[-?-?-AM-Q,W,WXaXgXg,h M )9iCJ?=-- 1 1 3Q7((; (8(8(?@DDFIZZQ ;;=#~~k; WdgtWt.;=NPW Y DGT!W'N%%dkk288:EU[[%&66*:: eDkFF3C)TUVW66w}Vcenpu7A5$,,XDG G (JN+?+?+AA AG Q488<< Q488<< Q488<< Q488<<xxzJ& 55y.g& phiop  p ps4R"B-R'' S0 R<<Sgt_maskr?rxyxyrxr|ctjd||}tj||d}t ||j d|z j S)aX Compute the instance segmentation loss for a single image. Args: gt_mask (torch.Tensor): Ground truth mask of shape (n, H, W), where n is the number of objects. pred (torch.Tensor): Predicted mask coefficients of shape (n, 32). proto (torch.Tensor): Prototype masks of shape (32, H, W). xyxy (torch.Tensor): Ground truth bounding boxes in xyxy format, normalized to [0, 1], of shape (n, 4). area (torch.Tensor): Area of each ground truth bounding box of shape (n,). Returns: (torch.Tensor): The calculated mask loss for a single image. Notes: The function uses the equation pred_mask = torch.einsum('in,nhw->ihw', pred, proto) to produce the predicted masks from the prototype masks and predicted mask coefficients. z in,nhw->ihwrr)r rprr)r einsumr$r%rr'r()rr?rrrx pred_maskr/s rsingle_mask_lossz#v8SegmentationLoss.single_mask_losssT(LLe< 11)WPVW$%**v*6=BBDDrrYrrrVrrrrc |j\} } } } d} ||gdz }t|dddfjd}|tj| | | | g|j z}t t|||||||D]\}}|\}}}}}}}|jrm||}| r*||dzjdddk(}|j}n||jd|k(|}| |j||||||||z } | |dzj|dzjzz } | |jz S) aF Calculate the loss for instance segmentation. Args: fg_mask (torch.Tensor): A binary tensor of shape (BS, N_anchors) indicating which anchors are positive. masks (torch.Tensor): Ground truth masks of shape (BS, H, W) if `overlap` is False, otherwise (BS, ?, H, W). target_gt_idx (torch.Tensor): Indexes of ground truth objects for each anchor of shape (BS, N_anchors). target_bboxes (torch.Tensor): Ground truth bounding boxes for each anchor of shape (BS, N_anchors, 4). batch_idx (torch.Tensor): Batch indices of shape (N_labels_in_batch, 1). proto (torch.Tensor): Prototype masks of shape (BS, 32, H, W). pred_masks (torch.Tensor): Predicted masks for each anchor of shape (BS, N_anchors, 32). imgsz (torch.Tensor): Size of the input image as a tensor of shape (2), i.e., (H, W). overlap (bool): Whether the masks in `masks` tensor overlap. Returns: (torch.Tensor): The calculated loss for instance segmentation. Notes: The batch loss can be computed for improved speed at higher memory usage. For example, pred_mask can be computed as follows: pred_mask = torch.einsum('in,nhw->ihw', pred, proto) # (i, 32) @ (32, 160, 160) -> (i, 160, 160) rr.rpNrr rJ) rcrprodr rPrR enumeratezipanyrOr&rr()rrYrrrVrrrrrrrrr/target_bboxes_normalizedmareamxyxyrsingle_i fg_mask_itarget_gt_idx_i pred_masks_iproto_imxyxy_imarea_imasks_imask_idxrs rrz.v8SegmentationLoss.calculate_segmentation_loss6sD %{{1ff$153F#F 23CG<AA!D)5<<QW8Xafamam+nn$S-UTY[`bg%hiKAx[c XI gwQX}}*95%(Q,)<)eZdZdZfdZdZedZdZxZ S) v8PoseLossrkct|||jdj|_t j |_|jddgk(}|jd}|r2tjtj|jn#tj||j|z }t||_y)zaInitializes v8PoseLoss with model, sets keypoint variables and declares a keypoint loss instance.rJrrr)rmN)rrr kpt_shaperrbce_poser from_numpyrrQrRonesrj keypoint_loss)rris_posenkptrmrs rrzv8PoseLoss.__init__{s R22,,. ..RG+~~a @G!!),// {{1T[[1$.uQx$>5E!Hy#(99`e-f`eZ\bggeAhnnQ6GRT.U`e-fhi#j#p#p \\A tww '$, ["))!Q2==? !))!Q2==? %%aA.99; !! U1X^^AB/ 5QTXT_T_`aTbb'3E4;;'L$ }!&&q) +&++B2 ))Ye (9(9"a(@%/RTUV//'**T[["9:TYZfTg/h&}}VQ7 9--4-044Q7&&}kB $$]NINN:r4cTXTbTb4cd BF--    ( ( *[-?-?-AM-Q,W,WXaXgXg,h M )9iCJ?=-- 1 1 3Q7((; (8(8(?@DDFIZZQ ;;= ] *M#~~k; Wdfs.? J DGT!Wk*--dkk:@@BHHJI f q )  f q ) #<qAEEG"KKXyq?QS\SbSbcdSe(f/8/?/?A z"A#IN3K:E a!6+"3"3A"6!66 7# "/!8!8!>q$>OD 'H.4ll2.>!.Cvf~*Y_`fYgimInH**8VXtLI~~b!Q& $ hv.>@P Q -''r) r3r4r5r6rrr7rrr8r9s@rrrxs*8986t>(rrceZdZdZdZy)v8ClassificationLossrkctjjj||dddz }|j }||fS)zDCompute the classification loss between predictions and true labels.rr(r@)r r functionalrbr)rrrr/ loss_itemss rrzv8ClassificationLoss.__call__ sDxx""00e PU0VY[[[[] ZrN)r3r4r5r6rrrrr'r' s 8 rr')r torch.nnrtorch.nn.functionalr*r$ultralytics.utils.metricsrultralytics.utils.opsrrrultralytics.utils.talrrr metricsr talr Modulerr;rErjr~rrr'rrrr4s /AANNBII*" "2"vryy"vJW299W"]6]6@b$b$JN(N(b  r