updates
This commit is contained in:
parent
209cc9e124
commit
06e88fec08
|
@ -487,7 +487,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru
|
||||||
# NMS methods https://github.com/ultralytics/yolov3/issues/679 'or', 'and', 'merge', 'vision', 'vision_batch'
|
# NMS methods https://github.com/ultralytics/yolov3/issues/679 'or', 'and', 'merge', 'vision', 'vision_batch'
|
||||||
|
|
||||||
# Box constraints
|
# Box constraints
|
||||||
min_wh, max_wh = 2, 10000 # (pixels) minimum and maximium box width and height
|
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximium box width and height
|
||||||
|
|
||||||
output = [None] * len(prediction)
|
output = [None] * len(prediction)
|
||||||
for image_i, pred in enumerate(prediction):
|
for image_i, pred in enumerate(prediction):
|
||||||
|
@ -524,8 +524,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru
|
||||||
|
|
||||||
# Batched NMS
|
# Batched NMS
|
||||||
if method == 'vision_batch':
|
if method == 'vision_batch':
|
||||||
i = torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres)
|
output[image_i] = pred[torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres)]
|
||||||
output[image_i] = pred[i]
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# All other NMS methods
|
# All other NMS methods
|
||||||
|
@ -541,8 +540,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru
|
||||||
dc = dc[:500] # limit to first 500 boxes: https://github.com/ultralytics/yolov3/issues/117
|
dc = dc[:500] # limit to first 500 boxes: https://github.com/ultralytics/yolov3/issues/117
|
||||||
|
|
||||||
if method == 'vision':
|
if method == 'vision':
|
||||||
i = torchvision.ops.boxes.nms(dc[:, :4], dc[:, 4], nms_thres)
|
det_max.append(dc[torchvision.ops.boxes.nms(dc[:, :4], dc[:, 4], nms_thres)])
|
||||||
det_max.append(dc[i])
|
|
||||||
|
|
||||||
elif method == 'or': # default
|
elif method == 'or': # default
|
||||||
# METHOD1
|
# METHOD1
|
||||||
|
|
Loading…
Reference in New Issue