Allow multi_label option for NMS with PyTorch Hub (#4728)
* Allow specifying multi_label option for NMS when using torch hub * Reformat Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
1 parent
deb434aefa
commit
1cad0ce2c7
1 file changed
+3
-1
+3
-1
@@ -278,6 +278,7 @@ class AutoShape(nn.Module):
|
||||
conf = 0.25 # NMS confidence threshold
|
||||
iou = 0.45 # NMS IoU threshold
|
||||
classes = None # (optional list) filter by class
|
||||
multi_label = False # NMS multiple labels per box
|
||||
max_det = 1000 # maximum number of detections per image
|
||||
|
||||
def __init__(self, model):
|
||||
@@ -337,7 +338,8 @@ class AutoShape(nn.Module):
|
||||
t.append(time_sync())
|
||||
|
||||
# Post-process
|
||||
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det) # NMS
|
||||
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes,
|
||||
multi_label=self.multi_label, max_det=self.max_det) # NMS
|
||||
for i in range(n):
|
||||
scale_coords(shape1, y[i][:, :4], shape0[i])
|
||||
|
||||
|
||||
Reference in new issue
Block a user