@@ -11,17 +11,23 @@ def nms(boxes, scores, iou_threshold):
11
11
IoU greater than iou_threshold with another (higher scoring)
12
12
box.
13
13
14
- Arguments:
15
- boxes (Tensor[N, 4]): boxes to perform NMS on. They
16
- are expected to be in (x1, y1, x2, y2) format
17
- scores (Tensor[N]): scores for each one of the boxes
18
- iou_threshold (float): discards all overlapping
19
- boxes with IoU < iou_threshold
20
-
21
- Returns:
22
- keep (Tensor): int64 tensor with the indices
23
- of the elements that have been kept
24
- by NMS, sorted in decreasing order of scores
14
+ Parameters
15
+ ----------
16
+ boxes : Tensor[N, 4])
17
+ boxes to perform NMS on. They
18
+ are expected to be in (x1, y1, x2, y2) format
19
+ scores : Tensor[N]
20
+ scores for each one of the boxes
21
+ iou_threshold : float
22
+ discards all overlapping
23
+ boxes with IoU < iou_threshold
24
+
25
+ Returns
26
+ -------
27
+ keep : Tensor
28
+ int64 tensor with the indices
29
+ of the elements that have been kept
30
+ by NMS, sorted in decreasing order of scores
25
31
"""
26
32
_C = _lazy_import ()
27
33
return _C .nms (boxes , scores , iou_threshold )
@@ -34,19 +40,25 @@ def batched_nms(boxes, scores, idxs, iou_threshold):
34
40
Each index value correspond to a category, and NMS
35
41
will not be applied between elements of different categories.
36
42
37
- Arguments:
38
- boxes (Tensor[N, 4]): boxes where NMS will be performed. They
39
- are expected to be in (x1, y1, x2, y2) format
40
- scores (Tensor[N]): scores for each one of the boxes
41
- idxs (Tensor[N]): indices of the categories for each
42
- one of the boxes.
43
- iou_threshold (float): discards all overlapping boxes
44
- with IoU < iou_threshold
45
-
46
- Returns:
47
- keep (Tensor): int64 tensor with the indices of
48
- the elements that have been kept by NMS, sorted
49
- in decreasing order of scores
43
+ Parameters
44
+ ----------
45
+ boxes : Tensor[N, 4]
46
+ boxes where NMS will be performed. They
47
+ are expected to be in (x1, y1, x2, y2) format
48
+ scores : Tensor[N]
49
+ scores for each one of the boxes
50
+ idxs : Tensor[N]
51
+ indices of the categories for each one of the boxes.
52
+ iou_threshold : float
53
+ discards all overlapping boxes
54
+ with IoU < iou_threshold
55
+
56
+ Returns
57
+ -------
58
+ keep : Tensor
59
+ int64 tensor with the indices of
60
+ the elements that have been kept by NMS, sorted
61
+ in decreasing order of scores
50
62
"""
51
63
if boxes .numel () == 0 :
52
64
return torch .empty ((0 ,), dtype = torch .int64 , device = boxes .device )
0 commit comments