Merge remote-tracking branch 'origin/master'
This commit is contained in:
commit
b9851222bf
|
|
@ -596,8 +596,6 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
|
||||||
Returns:
|
Returns:
|
||||||
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
||||||
"""
|
"""
|
||||||
if prediction.dtype is torch.float16:
|
|
||||||
prediction = prediction.float() # to FP32
|
|
||||||
|
|
||||||
nc = prediction[0].shape[1] - 5 # number of classes
|
nc = prediction[0].shape[1] - 5 # number of classes
|
||||||
xc = prediction[..., 4] > conf_thres # candidates
|
xc = prediction[..., 4] > conf_thres # candidates
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue