/j_ddlmZddlZddlmZddlmZddlmZm Z ddl m Z m Z m Z ddlmZGdd ejZGd d eZdd ZddZd d!dZd"dZ d#d$dZdS)%) annotationsN)LOGGER)bbox_iouprobiou) xywh2xyxyxywhr2xyxyxyxy xyxy2xywh) TORCH_1_11ceZdZdZddddgdddfdfd ZejdZdZdZ dZ dZ ddZ dZ ddZdZxZS)TaskAlignedAssigneraGA task-aligned assigner for object detection. This class assigns ground-truth (gt) objects to anchors based on the task-aligned metric, which combines both classification and localization information. Attributes: topk (int): The number of top candidates to consider. topk2 (int): Secondary topk value for additional filtering. num_classes (int): The number of object classes. alpha (float): The alpha parameter for the classification component of the task-aligned metric. beta (float): The beta parameter for the localization component of the task-aligned metric. stride (list): List of stride values for different feature levels. stride_val (int): The stride value used for select_candidates_in_gts. eps (float): A small value to prevent division by zero. P?g@) & .>Ntopkint num_classesalphafloatbetastridelistepsct||_|p||_||_||_||_||_t|jdkr |jdn |jd|_ ||_ dS)aInitialize a TaskAlignedAssigner object with customizable hyperparameters. Args: topk (int, optional): The number of top candidates to consider. num_classes (int, optional): The number of object classes. alpha (float, optional): The alpha parameter for the classification component of the task-aligned metric. beta (float, optional): The beta parameter for the localization component of the task-aligned metric. stride (list, optional): List of stride values for different feature levels. eps (float, optional): A small value to prevent division by zero. topk2 (int, optional): Secondary topk value for additional filtering. rrN) super__init__rtopk2rrrrlen stride_valr) selfrrrrrrr! __class__s Z/home/longshao/multi-rider-rag/.venv/lib/python3.11/site-packages/ultralytics/utils/tal.pyr zTaskAlignedAssigner.__init__s*  ]d &   ,/ ,<,z/TaskAlignedAssigner.forward..gs rrr1quuwwrrrr'c3BK|]}|VdSN)to)r,r-devices r& z.TaskAlignedAssigner.forward..is-::aQTT&\\::::::r'N)shapebs n_max_boxesr2torch full_liker zeros_like_forward RuntimeErrorstrlowerrwarningtuple) r$ pd_scores pd_bboxes anc_points gt_labels gt_bboxesmask_gte cpu_tensorsresultr2s @r&forwardzTaskAlignedAssigner.forward>sl,/!$$?1-!  q  & 143CDD ++ ++ 6!233 6!233   ==Iz9iY`aa a   #a&&,,..00XYYYrrIzS\^gip0qrrr & 4::::6:::::::::::  s%2C EA-E E E  Ec|||||||\}}} ||| |j|\} } }|||| | \} } }||z}|dd}| |zdd}||z||jzz dd}||z}| | || | fS)aCompute the task-aligned assignment. Args: pd_scores (torch.Tensor): Predicted classification scores with shape (bs, num_total_anchors, num_classes). pd_bboxes (torch.Tensor): Predicted bounding boxes with shape (bs, num_total_anchors, 4). anc_points (torch.Tensor): Anchor points with shape (num_total_anchors, 2). gt_labels (torch.Tensor): Ground truth labels with shape (bs, n_max_boxes, 1). gt_bboxes (torch.Tensor): Ground truth boxes with shape (bs, n_max_boxes, 4). mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (bs, n_max_boxes, 1). Returns: target_labels (torch.Tensor): Target labels with shape (bs, num_total_anchors). target_bboxes (torch.Tensor): Target bounding boxes with shape (bs, num_total_anchors, 4). target_scores (torch.Tensor): Target scores with shape (bs, num_total_anchors, num_classes). fg_mask (torch.Tensor): Foreground mask with shape (bs, num_total_anchors). target_gt_idx (torch.Tensor): Target ground truth indices with shape (bs, num_total_anchors). T)dimkeepdim) get_pos_maskselect_highest_overlapsr6 get_targetsamaxr unsqueezebool)r$r@rArBrCrDrEmask_pos align_metricoverlaps target_gt_idxfg_mask target_labels target_bboxes target_scorespos_align_metrics pos_overlapsnorm_align_metrics r&r:zTaskAlignedAssigner._forwardls$,0+<+< y)Y G, , (,,0+G+G h 0,, , ( w 7;6F6FyR[]jls6t6t3 }m  (--"d-CC 8+11b$1GG )L8io^g^noook47D,ITTU^__I z1---{7= 9K[\\\ yAAArr2%*YM]^^^ty!! L LA  % %b)AAAqqq!a!e)O*Dd K K K K!!,"2A666w}---r'ctj|jtj|jd}|||jzz}||}|d|j d|}| dtj |j d|j d|j ftj|j}| d|dd|dddddfdd|j } tj| dk|d}|||fS) a@Compute target labels, target bounding boxes, and target scores for the positive anchor points. Args: gt_labels (torch.Tensor): Ground truth labels of shape (b, max_num_obj, 1), where b is the batch size and max_num_obj is the maximum number of objects. gt_bboxes (torch.Tensor): Ground truth bounding boxes of shape (b, max_num_obj, 4). target_gt_idx (torch.Tensor): Indices of the assigned ground truth objects for positive anchor points, with shape (b, h*w), where h*w is the total number of anchor points. fg_mask (torch.Tensor): A boolean tensor of shape (b, h*w) indicating the positive (foreground) anchor points. Returns: target_labels (torch.Tensor): Target labels for positive anchor points with shape (b, h*w). target_bboxes (torch.Tensor): Target bounding boxes for positive anchor points with shape (b, h*w, 4). target_scores (torch.Tensor): Target scores for positive anchor points with shape (b, h*w, num_classes). )rlrjr2).NrKrrrirkN)r7ror5int64r2r6rnflattenrpr4r|rmrscatter_rSrepeatwhere) r$rCrDrXrY batch_indrZr[r\fg_scores_masks r&rQzTaskAlignedAssigner.get_targetssJ$LTWEK HXYYYZcd % D4D(DD !((0022=A "r9?2+>?? N  Q   #]%8%;T=M N+ '   q-"9"9""="=qAAA AAAt,33Aq$:JKK NQ$6 qII m]::r'ct|}|dddf|jdk}tj||ztj|j|j|j|dddf|dddf<t|}|j d}|j \}} } | ddd dd\} } tj |d| z | |dz fd  || |d} | d |S) aSelect positive anchor centers within ground truth bounding boxes. Args: xy_centers (torch.Tensor): Anchor center coordinates, shape (h*w, 2). gt_bboxes (torch.Tensor): Ground truth bounding boxes, shape (b, n_boxes, 4). mask_gt (torch.Tensor): Mask for valid ground truth boxes, shape (b, n_boxes, 1). eps (float, optional): Small value for numerical stability. Returns: (torch.Tensor): Boolean mask of positive anchors, shape (b, n_boxes, h*w). Notes: - b: batch size, n_boxes: number of ground truth boxes, h: height, w: width. - Bounding box format: [x_min, y_min, x_max, y_max]. .rkNrrirKrrL)r rr7rrTtensorr#rjr2rr4rpchunkcatamingt_)r$ xy_centersrDrErgt_bboxes_xywhwh_mask n_anchorsr5n_boxes_ltrb bbox_deltass r&rbz,TaskAlignedAssigner.select_candidates_in_gts!sF #9-- abb)DKN:"'+ w  $ $ & & L0D^Mb c c c 37 ## # sABBw n-- $Q' "GQAq))//155BiD!1B!6Z=M8M NTUVVV[[\^`girtvww ""&&s+++r'cL|d}|dkr|ddkd|d}|d}t j|j|j|j }| d|ddt j ||| }|d}|j |jkr~||z}t j||j ddj}t j|j|j|j } | d|d|| z}|d}|d} | ||fS)aSelect anchor boxes with highest IoU when assigned to multiple ground truths. Args: mask_pos (torch.Tensor): Positive mask, shape (b, n_max_boxes, h*w). overlaps (torch.Tensor): IoU overlaps, shape (b, n_max_boxes, h*w). n_max_boxes (int): Maximum number of ground truth boxes. align_metric (torch.Tensor): Alignment metric for selecting best matches. Returns: target_gt_idx (torch.Tensor): Indices of assigned ground truths, shape (b, h*w). fg_mask (torch.Tensor): Foreground mask, shape (b, h*w). mask_pos (torch.Tensor): Updated positive mask, shape (b, n_max_boxes, h*w). rNrrKriTrr)sumrrSreargmaxr7rmr4rjr2rrrr!rindices) r$rUrWr6rVrYmask_multi_gtsmax_overlaps_idxis_max_overlapstopk_idxrXs r&rPz+TaskAlignedAssigner.select_highest_overlaps@s,,r"" ;;==1  %//22Q6>>r;PRSSN'q11 #k(.W_WfgggO  $ $Q(8(B(B1(E(Eq I I I{>?HMMSSUUHll2&&G : " "'(2L$z, TXYYYa {8>PXP_```H   b"2C 8 8 8  Hll2&&G ++ gx//r') rrrrrrrrrrrrr0)r)__name__ __module__ __qualname____doc__r r7no_gradrIr:rOrcrrrdrQrbrP __classcell__)r%s@r&r r s$"{{>U]__++_+Z$Z$Z$ZL0004&&&B [ [ [....>';';';R,,,,>#0#0#0#0#0#0#0r'r ceZdZdZdZdZdS)RotatedTaskAlignedAssignerzSAssigns ground-truth objects to rotated bounding boxes using a task-aligned metric.cnt||ddS)z)Calculate IoU for rotated bounding boxes.rKr)rrqr|r}s r&rrz*RotatedTaskAlignedAssigner.iou_calculationis.y),,44R88??BBBr'cj|dddf|jdk}tj||ztj|j|j|j|dddf|dddf<t|}| dd\}}}} ||z } | |z } ||z } | | z d } | | z d }| | z d }| | z d }|dk|| kz|dkz||kzS) aSelect the positive anchor center in gt for rotated bounding boxes. Args: xy_centers (torch.Tensor): Anchor center coordinates with shape (h*w, 2). gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (b, n_boxes, 5). mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (b, n_boxes, 1). Returns: (torch.Tensor): Boolean mask of positive anchors with shape (b, n_boxes, h*w). .rkrrrirrNrrK) rr7rrTrr#rjr2r splitr)r$rrDrErcornersabrdabadapnorm_abnorm_ad ap_dot_ab ap_dot_ads r&rbz3RotatedTaskAlignedAssigner.select_candidates_in_gtsmsKC1H% A6#k w  $ $ & & L  HX Y Y Y c1Q3h   #qs(!++]]1"]-- 1a U U!^7--B-''7--B-''"WMMbM)) "WMMbM)) Q9#78INKy\cOcddr'N)rrrrrrrbr*r'r&rrfs@]]CCCeeeeer'r?c ^gg}}|J|dj|dj}}tt|D]F}||}t |t r||jddn5t||dt||df\} } tj | |||z} tj | |||z} trtj | | dntj | | \} } | tj | | fddd| tj| | zdf||| Htj|tj|fS) zGenerate anchors from features.Nrrkr)rlr2rjij)indexingrKri)rjr2rr" isinstancerr4rr7ror meshgridappendstackrpfullr) featsstridesgrid_cell_offset anchor_points stride_tensorrjr2irhwsxsys r& make_anchorsrs#%r=M    !HNE!HO6E 3u::  YY%/t%<%<fuQx~abb!!3uQxPQ{CSCSUXY^_`YabcYdUeUeBf1 \ae < < )B 88A38  DAq1urzH1urzH1urzH1urzH 9h(H=3 G G GD{{1gn-- Kr')r)TrKr0)rrrrrrrr)rK)rKN) r[rrrrrrLrrr) __future__rr7torch.nnnnrrrropsrr r torch_utilsr Moduler rrrrrrr*r'r&rsw#""""" &&&&&&&&5555555555######U0U0U0U0U0")U0U0U0p &e&e&e&e&e!4&e&e&eR > > > > ( ( ( (----2 $$$$$$$r'