o W he@spddlZddlmZddlmmZddlmZddlm Z m Z m Z ddl m Z mZmZmZmZddlmZddlmZmZddlmZGd d d ejZGd d d ejZGd ddejZGdddejZGdddeZGdddejZGdddZ Gddde Z!Gddde Z"GdddZ#Gddde Z$Gdd d Z%Gd!d"d"Z&Gd#d$d$e&Z'dS)%N) OKS_SIGMA) crop_mask xywh2xyxy xyxy2xywh)RotatedTaskAlignedAssignerTaskAlignedAssigner dist2bbox dist2rbox make_anchors)autocast)bbox_iouprobiou) bbox2distc*eZdZdZdfdd ZddZZS) VarifocalLossa Varifocal loss by Zhang et al. https://arxiv.org/abs/2008.13367. Args: gamma (float): The focusing parameter that controls how much the loss focuses on hard-to-classify examples. alpha (float): The balancing factor used to address class imbalance. @?ct||_||_dS)z#Initialize the VarifocalLoss class.Nsuper__init__gammaalphaselfrr __class__J/var/www/vscode/kcb/lib/python3.10/site-packages/ultralytics/utils/loss.pyr  zVarifocalLoss.__init__cCs||j||jd|||}tddtj||dd|d }Wd|S1s7wY|S)zrH)r&rOr`r+rsrNexpr*)r pred_kptsgt_kptskpt_maskareadkpt_loss_factorerrrr1s,  $zKeypointLoss.forwardrXr2rrrrrrrorrc@s2eZdZdZd ddZddZddZd d Zd S) v8DetectionLosszJCriterion class for computing training losses for YOLOv8 object detection. cCst|j}|j}|jd}tjdd|_||_|j |_ |j |_ |j |j d|_ |j |_ ||_|j dk|_ t||j ddd|_t|j ||_tj|j tj|d |_d S) zVInitialize v8DetectionLoss with model parameters and task-aligned assignment settings.rHr"r#rqr ?@topk num_classesrbeta)dtypercN)next parametersrcargsmodelnnBCEWithLogitsLossbcehypstridencrEnouse_dflrassignerrYrb bbox_lossr`aranger)proj)rrtal_topkrchmrrrrs  zv8DetectionLoss.__init__c Cs|j\}}|dkrtj|d|d|jd}|S|dddf}|jdd\}} | jtjd} tj|| |d|jd}t|D]} || k} | } rZ|| ddf|| d| f<q@t |ddd f ||ddd f<|S) zJPreprocess targets by converting to tensor format and scaling coordinates.rr rcNT return_countsr.) rOr`zerosrcuniquerbint32maxranger+rmul_) rtargets batch_size scale_tensornlneouti_countsjmatchesnrrr preprocesss    &zv8DetectionLoss.preprocesscCsJ|jr|j\}}}|||d|dd|j|j}t||ddS)zUDecode predicted object bounding box coordinates from anchor points and distribution.rqF)r\) rrOrNsoftmaxmatmulrtyperr)rrfrPbacrrr bbox_decodes *zv8DetectionLoss.bbox_decodec s,tjdjd}t|tr|dn|tfddDdjdjfd\}}| ddd }| ddd }|j }|j d}tj dj dd j|d jd}tjd \} } t|d d d|dd d|dfd} j| j||gdd} | dd\} } | jdddd}| |}||| | j | | | | |\}}}}}t|d}|||||d<|r|| }||| ||||\|d<|d<|djj9<|djj9<|djj 9<|||fS)LCalculate the sum of the loss for box, cls and dfl multiplied by batch size.rrr c&g|]}|djdjdqSrrHrNrOr.0xifeatsrrr &z,v8DetectionLoss.__call__..rurqrNrcrr batch_idxrHclsbboxesr rr rrr rqTrIr^)!r`rrc isinstancetuplecatsplitrErpermute contiguousrrOrarr rNrrbr+gt_rrdetachr%rrrrrboxrdfl)rpredsbatchr0 pred_distri pred_scoresrrimgszrf stride_tensorr gt_labels gt_bboxesmask_gtrerrgrhrjrirrrrVsF *.    zv8DetectionLoss.__call__N)r)r3r4r5r6rrrrVrrrrrs   rcseZdZdZfddZddZedejdejdejd ejd ejd ejf d d Z dejdejdejdejdejdejdejdejde d ejfddZ Z S)v8SegmentationLosszFCriterion class for computing training losses for YOLOv8 segmentation.cst||jj|_dS)zWInitialize the v8SegmentationLoss class with model parameters and mask overlap setting.N)rrr overlap_maskoverlaprrrrrrs zv8SegmentationLoss.__init__c stjdjd}t|dkr|n|d\}}|j\}}}} tfddDdjdjfd\} } | ddd } | ddd } | ddd }| j } tj djdd j| d j d} tj d \}}z=|d d d}t||dd d|dfd}j|j|| gdd}|dd\}}|jdddd}Wnty}ztd|d }~ww|| }| |||j |||||\}}}}}t|d}| || ||d<|rK| |||||||\|d<|d<|dj}t|jdd || fkr:t j!|d || fddd}"|||||||| j# |d<n|d|d|d7<|dj$j%9<|dj$j%9<|dj$j&9<|dj$j'9<|||fS)zFCalculate and return the combined loss for detection and segmentation.rqrrr crrrrrrrrrz/v8SegmentationLoss.__call__..rurNrrrrHrrrrrTrIr^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=yolo11n-seg.pt data=coco8.yaml'. Verify your dataset is a correctly formatted 'segment' dataset using 'data=coco8-seg.yaml' as an example. See https://docs.ultralytics.com/datasets/segment/ for help.masksnearest)mode)(r`rrclenrOrrrErrrrrarr rNrrbr+r RuntimeError TypeErrorrrrr%rrrrr)rr' interpolatecalculate_segmentation_lossrrrrr)rrrr0 pred_masksprotorrmask_hmask_wrrrrrfrrrrrrrrergrhrj target_gt_idxrirrrrrVsv*"       $zv8SegmentationLoss.__call__gt_maskr<rxyxyr}rCcCs8td||}tj||dd}t||jdd|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->ihwr"r#)r rurw)r`einsumr'r(rr*r+)rr<rrr} pred_maskr0rrrsingle_mask_lossisz#v8SegmentationLoss.single_mask_lossrjrrrgrrrrc  Cs"|j\} } } } d} ||gd}t|dddfd}|tj| | | | g|jd}tt|||||||D]R\}}|\}}}}}}}|r|||}| r_||d dddk}| }n || d|k|}| | ||||||||7} q8| |d |d 7} q8| | 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.ruNrr rH) rOrprodr`rarc enumeratezipanyrNr)rr+)rrjrrrgrrrrrrrrr0target_bboxes_normalizedmareamxyxyrsingle_i fg_mask_itarget_gt_idx_i pred_masks_iproto_imxyxy_imarea_imasks_imask_idxrrrrrs$"   z.v8SegmentationLoss.calculate_segmentation_loss) r3r4r5r6rrV staticmethodr`Tensorrboolrr7rrrrrsN S    rcs<eZdZdZfddZddZeddZdd ZZ S) v8PoseLosszICriterion class for computing training losses for YOLOv8 pose estimation.csvt||jdj|_t|_|jddgk}|jd}|r)tt  |j n tj ||j d|}t |d|_dS)zQInitialize v8PoseLoss with model parameters and keypoint-specific loss functions.rHrrr)rsN)rrr kpt_shaperrbce_poser` from_numpyrrbrconesrr keypoint_loss)rris_posenkptrsrrrrs   *zv8PoseLoss.__init__c stjdjd}t|dtr|n|d\}tfddDdjdjfd\}}| ddd }| ddd }| ddd }|j }tj dj dd j|d jd}tjd \} } |j d} |d d d} t| |dd d|dfd} j| j| |gdd} | dd\}}|jdddd}| |}| |j| d gjR}||| |j | | |||\}}}}}t|d}|||||d<|r@|| }||| ||||\|d<|d<|dj }|d|d9<|d|d9<!|||| | ||\|d<|d<|dj"j#9<|dj"j$9<|dj"j%9<|dj"j&9<|dj"j'9<|| |fS)z;Calculate the total loss and detach it for pose estimation.rrrr crrrrrrrrrz'v8PoseLoss.__call__..rurqNrrrrHrrrrrTrIr^r keypointsrtrv)(r`rrcrlistrrrErrrrrarOrr rNrrbr+rr kpts_decoderrrr%rrrrr)clonecalculate_keypoints_lossrrposekobjrr)rrrr0rzrrrrrfrrrrrrrrerrgrhrjrrirrrrrVs\* "     zv8PoseLoss.__call__cCsh|}|dddfd9<|d|dddgfd7<|d|ddd gfd7<|S) z0Decode predicted keypoints to image coordinates..Nrurrtrrrvr )r)rfrzyrrrrs ""zv8PoseLoss.kpts_decodec Cs||}t|}tj|ddd} tj|| |jd|jdf|jd} t|D]} ||| k} | | | d| jdf<q)| d d} | d| dd|jd|jd}|d ddf| dddd<d}d}| r||}t||ddddfjddd }||}|jdd kr|d dknt|d d}|||||}|jdd kr||d |}||fS)a Calculate the keypoints loss for the model. This function calculates the keypoints loss and keypoints object loss for a given batch. The keypoints loss is based on the difference between the predicted keypoints and ground truth keypoints. The keypoints object loss is a binary classification loss that classifies whether a keypoint is present or not. Args: masks (torch.Tensor): Binary mask tensor indicating object presence, shape (BS, N_anchors). target_gt_idx (torch.Tensor): Index tensor mapping anchors to ground truth objects, shape (BS, N_anchors). keypoints (torch.Tensor): Ground truth keypoints, shape (N_kpts_in_batch, N_kpts_per_object, kpts_dim). batch_idx (torch.Tensor): Batch index tensor for keypoints, shape (N_kpts_in_batch, 1). stride_tensor (torch.Tensor): Stride tensor for anchors, shape (N_anchors, 1). target_bboxes (torch.Tensor): Ground truth boxes in (x1, y1, x2, y2) format, shape (BS, N_anchors, 4). pred_kpts (torch.Tensor): Predicted keypoints, shape (BS, N_anchors, N_kpts_per_object, kpts_dim). Returns: kpts_loss (torch.Tensor): The keypoints loss. kpts_obj_loss (torch.Tensor): The keypoints object loss. Trr rurNrrH.rIr).rurt)flattenrr`rrrrOrcrr_gatherexpandrNrrr full_likerrr))rrrrrrrgrzrmax_kptsbatched_keypointsr keypoints_itarget_gt_idx_expandedselected_keypoints kpts_loss kpts_obj_lossgt_kptr}pred_kptr|rrrrs2  $&*z#v8PoseLoss.calculate_keypoints_loss) r3r4r5r6rrVr rrr7rrrrr s  B r c@seZdZdZddZdS)v8ClassificationLosszACriterion class for computing training losses for classification.cCs>t|ttfr |dn|}tj||ddd}|}||fS)zDCompute the classification loss between predictions and true labels.r rr*r#)rrrr'rMr)rrrr0 loss_itemsrrrrVcszv8ClassificationLoss.__call__N)r3r4r5r6rVrrrrr*`s r*cs8eZdZdZfddZddZddZdd ZZS) v8OBBLosszdCalculates losses for object detection, classification, and box distribution in rotated YOLO models.cs8t|td|jddd|_t|j|j|_ dS)z^Initialize v8OBBLoss with model, assigner, and rotated bbox loss; model must be de-paralleled.rrrrN) rrrrrrprErbrcrrrrrrns zv8OBBLoss.__init__c Cs|jddkrtj|dd|jd}|S|dddf}|jdd\}}|jtjd}tj||d|jd}t|D]4}||k} | } rn|| ddf} | d dd f |tj || d df| gd d ||d| f<q:|S)z7Preprocess targets for oriented bounding box detection.rrNTrrru.rqr rHrw) rOr`rrcrrbrrrr+rr) rrrrrrrrrrrrrrrrts   *zv8OBBLoss.preprocessc stjdjd}t|dtr|n|d\}|jd}tfddDdjdj fd\}}| ddd }| ddd }| ddd }|j }tj djdd j|d jd} tjd \} } zh|d d d} t| |dd d|dd dfd} | d d df| d| d d df| d}}| |dk|dk@} j| j|| gdd} | dd\}}|jdddd}Wnty}ztd|d }~ww| ||}|}|dd df| 9<|||j | | |||\}}}}}t|d}|||||d<|rW|dd df| < ||| ||||\|d<|d<n |d|d7<|dj!j"9<|dj!j#9<|dj!j$9<|||fS)zBCalculate and return the loss for oriented bounding box detection.rrrr crrrrrrrrrz&v8OBBLoss.__call__..rurqNrrrrHrrrrr)r rTrIr^uhERROR ❌ OBB dataset incorrectly formatted or not a OBB dataset. This error can occur when incorrectly training a 'OBB' model on a 'detect' dataset, i.e. 'yolo train model=yolo11n-obb.pt data=dota8.yaml'. Verify your dataset is a correctly formatted 'OBB' dataset using 'data=dota8.yaml' as an example. See https://docs.ultralytics.com/datasets/obb/ for help..)%r`rrcrrrOrrrErrrrrarr rNitemrrbr+rrrrrrrr%rrrrrrrr)rrrr0 pred_anglerrrrrrfrrrrwrhrrrrrebboxes_for_assignerrrgrhrjrirrrrVsf **:      zv8OBBLoss.__call__cCsV|jr|j\}}}|||d|dd|j|j}tj t ||||fddS)a Decode predicted object bounding box coordinates from anchor points and distribution. Args: anchor_points (torch.Tensor): Anchor points, (h*w, 2). pred_dist (torch.Tensor): Predicted rotated distance, (bs, h*w, 4). pred_angle (torch.Tensor): Predicted angle, (bs, h*w, 1). Returns: (torch.Tensor): Predicted rotated bounding boxes with angles, (bs, h*w, 5). rqrrHrw) rrOrNrrrrrr`rr )rrfrPr/rrrrrrrs *zv8OBBLoss.bbox_decode) r3r4r5r6rrrVrr7rrrrr,ks  Hr,c@ eZdZdZddZddZdS) E2EDetectLosszGCriterion class for computing training losses for end-to-end detection.cCs t|dd|_t|dd|_dS)zcInitialize E2EDetectLoss with one-to-many and one-to-one detection losses using the provided model.r)rr N)rone2manyone2onerrrrrszE2EDetectLoss.__init__cCs^t|tr |dn|}|d}|||}|d}|||}|d|d|d|dfS)rr r5r6r)rrr5r6)rrrr5 loss_one2manyr6 loss_one2onerrrrVs    zE2EDetectLoss.__call__Nr3r4r5r6rrVrrrrr4s r4c@s(eZdZdZddZddZddZdS) TVPDetectLosszOCriterion class for computing training losses for text-visual prompt detection.cCs,t||_|jj|_|jj|_|jj|_dS)z^Initialize TVPDetectLoss with task-prompt and visual-prompt criteria using the provided model.N)r vp_criterionrori_ncrori_norE ori_reg_maxrrrrrs   zTVPDetectLoss.__init__cCst|tr |dn|}|j|jjksJ|jd|j|djdkr3tjd|jj dd}|| fS| |}|||}|dd}||dfS)z4Calculate the loss for text-visual prompt detection.r rqrrTrc requires_grad) rrr>r;rEr<rOr`rrcr_get_vp_features)rrrrr0vp_featsvp_lossbox_lossrrrrVs     zTVPDetectLoss.__call__csb|djdjdjj_jjdj_jj_ddfdd|DDS)z5Extract visual-prompt features from the model output.rr rqcSs$g|]\}}}tj||fddqS)r rw)r`r)rrrcls_vprrrrsz2TVPDetectLoss._get_vp_features..cs(g|]}|jjdjfddqS)rqr rw)rr>r<rrvncrrrs() rOr>r<r;rrErrr)rrrrFrrA s zTVPDetectLoss._get_vp_featuresN)r3r4r5r6rrVrArrrrr:s  r:c@r3)TVPSegmentLosszRCriterion class for computing training losses for text-visual prompt segmentation.cCst||_dS)z_Initialize TVPSegmentLoss with task-prompt and visual-prompt criteria using the provided model.N)rr;rrrrrszTVPSegmentLoss.__init__c Cst|dkr|n|d\}}}|jj|jjksJ|jjd|jj|djdkr:tjd|jjdd}|| fS| |}||||f|}|dd} | |dfS)z7Calculate the loss for text-visual prompt segmentation.rr rqrTr?ru) r tp_criterionrEr;rrOr`rrcrrA) rrrrrrr0rBrCcls_lossrrrrVs"    zTVPSegmentLoss.__call__Nr9rrrrrHs rH)(r`torch.nnrtorch.nn.functional functionalr'ultralytics.utils.metricsrultralytics.utils.opsrrrultralytics.utils.talrrrr r ultralytics.utils.torch_utilsr metricsr rtalrModulerr8rArYrprrrrr r*r,r4r:rHrrrrs2     j8 t'