Limit detections without explicit if condition (#10502)
* limit detections without explicit if condition Signed-off-by: Yonghye Kwon <developer.0hye@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup indexing code for limit detections Signed-off-by: Yonghye Kwon <developer.0hye@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
1 file changed
+1
-2
+1
-2
@@ -978,8 +978,7 @@ def non_max_suppression(
|
|||||||
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
||||||
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
||||||
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
||||||
if i.shape[0] > max_det: # limit detections
|
i = i[:max_det] # limit detections
|
||||||
i = i[:max_det]
|
|
||||||
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
||||||
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
||||||
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
||||||
|
|||||||
Reference in new issue
Block a user