Fi)8ddlZddlZddlZddlmZddlmZmZddlm Z  dde de de de d e d e d e d e d e de de de fdZ GddZy)N)LOGGER) batch_probioubox_iou) xywh2xyxy conf_thres iou_thresagnostic multi_labelmax_detnc max_time_imgmax_nmsmax_whrotatedend2end return_idxsc  d|cxkrdks nJd|dd|cxkrdks nJd|dt|ttfr|d}|!tj||j }|j dd k(s| rP|Dcgc]}||ddd f|kDd|}}|-|Dcgc]"}||ddd d f|k(jd$}}|S|j d}|xs|j dd z }|j d|z d z }d |z}|ddd |fjd|kD}tj|j d|j j|dd }d | |zz}||dkDz}|jdd}| st|ddd f|ddd f<tj}tjdd |zf|j g|z}tjd|j g|z}tt!||D]D\}\}}||}||}|r||}|rt#||r| s||}tjt#|||zd zf|j }t|dddd f|dddd f<d|t%t#||dddfj'd zf<tj(||fd}|j ds|j+d ||fd\}} }!|ratj,| |kD\}"}#tj(||"||"d |#zdf|#dddfj/|!|"fd}|rb||"}n\| j1dd\}$}#|$j3d|kD}tj(||$|#j/|!fd|}|r||}|)|ddd d f|k(jd}||}|r||}|j d}%|%s|%| kDr(|ddd fj5dd| }||}|r||}|ddd d f|rdn| z}&|ddd f}'| rVtj(|ddddf|&z|dddd f|ddddffd}(t6j9|(|'|t:}"nY|dddd f|&z}(dt<j>vr"ddl })|)jBjE|(|'|}"nt6jE|(|'|}"|"d|}"||"||<|r||"j3d||<tj|z |kDs+tGjHd|ddn|r||fS|Scc}wcc}w)adPerform non-maximum suppression (NMS) on prediction results. Applies NMS to filter overlapping bounding boxes based on confidence and IoU thresholds. Supports multiple detection formats including standard boxes, rotated boxes, and masks. Args: prediction (torch.Tensor): Predictions with shape (batch_size, num_classes + 4 + num_masks, num_boxes) containing boxes, classes, and optional masks. conf_thres (float): Confidence threshold for filtering detections. Valid values are between 0.0 and 1.0. iou_thres (float): IoU threshold for NMS filtering. Valid values are between 0.0 and 1.0. classes (list[int], optional): List of class indices to consider. If None, all classes are considered. agnostic (bool): Whether to perform class-agnostic NMS. multi_label (bool): Whether each box can have multiple labels. labels (list[torch.Tensor]): A priori labels for each image. max_det (int): Maximum number of detections to keep per image. nc (int): Number of classes. Indices after this are considered masks. max_time_img (float): Maximum time in seconds for processing one image. max_nms (int): Maximum number of boxes for NMS. max_wh (int): Maximum box width and height in pixels. rotated (bool): Whether to handle Oriented Bounding Boxes (OBB). end2end (bool): Whether the model is end-to-end and doesn't require NMS. return_idxs (bool): Whether to return the indices of kept detections. Returns: (list[torch.Tensor] | tuple[list[torch.Tensor], list[torch.Tensor]]): List of detections per image with shape (num_boxes, 6 + num_masks) containing (x1, y1, x2, y2, confidence, class, mask1, mask2, ...). If return_idxs=True, returns a tuple of (output, keepi) where keepi contains indices of kept detections. rzInvalid Confidence threshold z&, valid values are between 0.0 and 1.0z Invalid IoU Ndevice).Ng@.)rrg?T)keepdim descending)dim)iou_func torchvisionzNMS time limit z.3fz s exceeded)% isinstancelisttupletorchtensorrshapeanyamaxarangeexpand transposertimezeros enumerateziplenrangelongcatsplitwherefloatmaxviewargsortTorchNMSfast_nmsrsysmodulesr"opsnmsrwarning)* predictionrrclassesr r labelsr r r rrrrrpredoutputbsextramixcxinds time_limittkeepixixxkfiltlbvboxclsmaskijconfncscoresboxesr"s* W/var/www/pod-logistic/pod-ai/venv/lib/python3.12/site-packages/ultralytics/utils/nms.pynon_max_suppressionra s\ a s#@ Lr!ss  Q `,yk9_ `` *tUm,] ,,wz/@/@Aq GFPQd$tAqDzJ./9QQ  IOPdDAaCLG388;<PFP   ! B (   #a'B   Q " $q (E RB AqtG  ! !! $z 1B LL))"-j6G6G H O OPRTV WXa bE|b((J26K%%b"-J ' 37(;< 37 Akk1a%i.1B1BC Dr IF [[ (9(9 : ;b @E Z!78I GQ"v dG DB c&*ogB SWb5j1n5ahhGA AqsF,Aa!eH58AeCGnbAhmmo11 2 1a&!$Awwqz !R3S$ ;;sZ/0DAq 3q61QAt^#4a4j6F6F6H$q'RTUVAUggag.GD!99R=:-D 3aggi6:4@AX  a1fI(--a0D$AX GGAJ  w;QT7??d?3HW=D$AX a1fIhF 31a4 IIqBQBx!|Qq!A#vY!RS& BKE!!%]!SAa!eHqLE +"OO''vyALL : hwKqTr 1 2E"I IIK!Oz ) NN_Z,>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]]) >>> scores = torch.tensor([0.9, 0.8]) >>> keep = TorchNMS.nms(boxes, scores, 0.5) Tr_r^ iou_thresholduse_triu exit_earlyreturnc^|jdk(r2|r0tjdtj|jStj |d}||}|||}|rP|j d}tj||k\jddkjd}||S|jd} tj| |j jddjd| } tj| |j jddj| d} | | k} || z}||} d| ||k\jddk<| ||<tj| | jdj}||S) aFast-NMS implementation from https://arxiv.org/pdf/1904.02689 using upper triangular matrix operations. Args: boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format. scores (torch.Tensor): Confidence scores with shape (N,). iou_threshold (float): IoU threshold for suppression. use_triu (bool): Whether to use torch.triu operator for upper triangular matrix operations. iou_func (callable): Function to compute IoU between boxes. exit_early (bool): Whether to exit early if there are no boxes. Returns: (torch.Tensor): Indices of boxes to keep after NMS. Examples: Apply NMS to a set of boxes >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]]) >>> scores = torch.tensor([0.9, 0.8]) >>> keep = TorchNMS.fast_nms(boxes, scores, 0.5) rrdtyperTrr)diagonalrr)numelr&emptyint64rr;triu_nonzerosumsqueeze_r(r+r:r,topkindices)r_r^rcrdr!re sorted_idxiouspickr\row_idxcol_idx upper_maskscores_s r`r=zTorchNMS.fast_nmss8 ;;=A *;;t5;;u||L L]]6d; j!u% ::q:)D==$-"7!?Gt},11!49: ;!(F: ::gw}}Q'78@@D$c|jdk(r0tjdtj|jS|j d\}}}}||z ||z z}|j dd}tj|jtj|j} d} |jdkDr |d} | | | <| dz } |jdk(rn|dd} tj|| || } tj|| || }tj|| || }tj|| || }|| z jd}||z jd}||z}|jdk(r| }||| || z|z z }| ||k}|jdkDr | d| S) aOptimized NMS with early termination that matches torchvision behavior exactly. Args: boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format. scores (torch.Tensor): Confidence scores with shape (N,). iou_threshold (float): IoU threshold for suppression. Returns: (torch.Tensor): Indices of boxes to keep after NMS. Examples: Apply NMS to a set of boxes >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]]) >>> scores = torch.tensor([0.9, 0.8]) >>> keep = TorchNMS.nms(boxes, scores, 0.5) rrhrirTrN)min) rlr&rmrnrunbindr;r/maximumminimumclamp_rq)r_r^rcx1y1x2y2areasorderkeepkeep_idxrYrestxx1yy1xx2yy2whinterious r`rAz TorchNMS.nmss$ ;;=A ;;t5;;u||L LaBBbR"W%qT2{{5;;= ELLQkkmaaADN MH{{}!9D--1r$x0C--1r$x0C--1r$x0C--1r$x0Cs""q")As""q")AEEyy{a58eDk1E9:C -.E3kkma6IXr|idxs use_fast_nmscN|jdk(r0tjdtj|jS|j }|j ||dzz}||dddfz}|rtj|||Stj|||S)a.Batched NMS for class-aware suppression. Args: boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format. scores (torch.Tensor): Confidence scores with shape (N,). idxs (torch.Tensor): Class indices with shape (N,). iou_threshold (float): IoU threshold for suppression. use_fast_nms (bool): Whether to use the Fast-NMS implementation. Returns: (torch.Tensor): Indices of boxes to keep after NMS. Examples: Apply batched NMS across multiple classes >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]]) >>> scores = torch.tensor([0.9, 0.8]) >>> idxs = torch.tensor([0, 1]) >>> keep = TorchNMS.batched_nms(boxes, scores, idxs, 0.5) rrhrirN) rlr&rmrnrr9tor<r=rA)r_r^rrcrmax_coordinateoffsets boxes_for_nmss r` batched_nmszTorchNMS.batched_nms+s6 ;;=A ;;t5;;u||L L''%.NQ$674 00    mV] C mV]C r|N)F) __name__ __module__ __qualname____doc__ staticmethodrr&Tensorr8boolr=rArr|r`r<r<s "  1 ||1  1 1  1  1  1 1 f95<<99e9PUP\P\99v # & ||&  & ll&  &  &  & & r|r<)g?g?NFFri,rg?i0uiFFF)r>r.r&ultralytics.utilsrultralytics.utils.metricsrrultralytics.utils.opsrr8rintrar<rr|r`rs $<+   Y6Y6Y6  Y6  Y6Y6 Y6Y6Y6 Y6Y6Y6Y6xi i r|