Commit
•
3f74cd9
1
Parent(s):
13a1c72
Parameterize max_det + inference default at 1000 (#3215)
Browse files* Added max_det parameters in various places
* 120 character line
* PEP8
* 120 character line
* Update inference default to 1000 instances
* Update inference default to 1000 instances
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
- detect.py +3 -1
- models/common.py +4 -2
- utils/general.py +1 -2
detect.py
CHANGED
@@ -68,7 +68,8 @@ def detect(opt):
|
|
68 |
pred = model(img, augment=opt.augment)[0]
|
69 |
|
70 |
# Apply NMS
|
71 |
-
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres,
|
|
|
72 |
t2 = time_synchronized()
|
73 |
|
74 |
# Apply Classifier
|
@@ -153,6 +154,7 @@ if __name__ == '__main__':
|
|
153 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
154 |
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|
155 |
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
|
|
|
156 |
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
157 |
parser.add_argument('--view-img', action='store_true', help='display results')
|
158 |
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
|
|
68 |
pred = model(img, augment=opt.augment)[0]
|
69 |
|
70 |
# Apply NMS
|
71 |
+
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms,
|
72 |
+
max_det=opt.max_det)
|
73 |
t2 = time_synchronized()
|
74 |
|
75 |
# Apply Classifier
|
|
|
154 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
155 |
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|
156 |
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
|
157 |
+
parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
|
158 |
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
159 |
parser.add_argument('--view-img', action='store_true', help='display results')
|
160 |
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
models/common.py
CHANGED
@@ -215,12 +215,13 @@ class NMS(nn.Module):
|
|
215 |
conf = 0.25 # confidence threshold
|
216 |
iou = 0.45 # IoU threshold
|
217 |
classes = None # (optional list) filter by class
|
|
|
218 |
|
219 |
def __init__(self):
|
220 |
super(NMS, self).__init__()
|
221 |
|
222 |
def forward(self, x):
|
223 |
-
return non_max_suppression(x[0],
|
224 |
|
225 |
|
226 |
class AutoShape(nn.Module):
|
@@ -228,6 +229,7 @@ class AutoShape(nn.Module):
|
|
228 |
conf = 0.25 # NMS confidence threshold
|
229 |
iou = 0.45 # NMS IoU threshold
|
230 |
classes = None # (optional list) filter by class
|
|
|
231 |
|
232 |
def __init__(self, model):
|
233 |
super(AutoShape, self).__init__()
|
@@ -285,7 +287,7 @@ class AutoShape(nn.Module):
|
|
285 |
t.append(time_synchronized())
|
286 |
|
287 |
# Post-process
|
288 |
-
y = non_max_suppression(y,
|
289 |
for i in range(n):
|
290 |
scale_coords(shape1, y[i][:, :4], shape0[i])
|
291 |
|
|
|
215 |
conf = 0.25 # confidence threshold
|
216 |
iou = 0.45 # IoU threshold
|
217 |
classes = None # (optional list) filter by class
|
218 |
+
max_det = 1000 # maximum number of detections per image
|
219 |
|
220 |
def __init__(self):
|
221 |
super(NMS, self).__init__()
|
222 |
|
223 |
def forward(self, x):
|
224 |
+
return non_max_suppression(x[0], self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det)
|
225 |
|
226 |
|
227 |
class AutoShape(nn.Module):
|
|
|
229 |
conf = 0.25 # NMS confidence threshold
|
230 |
iou = 0.45 # NMS IoU threshold
|
231 |
classes = None # (optional list) filter by class
|
232 |
+
max_det = 1000 # maximum number of detections per image
|
233 |
|
234 |
def __init__(self, model):
|
235 |
super(AutoShape, self).__init__()
|
|
|
287 |
t.append(time_synchronized())
|
288 |
|
289 |
# Post-process
|
290 |
+
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det) # NMS
|
291 |
for i in range(n):
|
292 |
scale_coords(shape1, y[i][:, :4], shape0[i])
|
293 |
|
utils/general.py
CHANGED
@@ -482,7 +482,7 @@ def wh_iou(wh1, wh2):
|
|
482 |
|
483 |
|
484 |
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
485 |
-
labels=()):
|
486 |
"""Runs Non-Maximum Suppression (NMS) on inference results
|
487 |
|
488 |
Returns:
|
@@ -498,7 +498,6 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
|
|
498 |
|
499 |
# Settings
|
500 |
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
501 |
-
max_det = 300 # maximum number of detections per image
|
502 |
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
503 |
time_limit = 10.0 # seconds to quit after
|
504 |
redundant = True # require redundant detections
|
|
|
482 |
|
483 |
|
484 |
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
485 |
+
labels=(), max_det=300):
|
486 |
"""Runs Non-Maximum Suppression (NMS) on inference results
|
487 |
|
488 |
Returns:
|
|
|
498 |
|
499 |
# Settings
|
500 |
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
|
|
501 |
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
502 |
time_limit = 10.0 # seconds to quit after
|
503 |
redundant = True # require redundant detections
|