Unverified 提交 68e6ab66 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Hub device mismatch bug fix (#1619)

上级 791dadb5
...@@ -265,7 +265,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None, ...@@ -265,7 +265,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None,
detections with shape: nx6 (x1, y1, x2, y2, conf, cls) detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
""" """
nc = prediction[0].shape[1] - 5 # number of classes nc = prediction.shape[2] - 5 # number of classes
xc = prediction[..., 4] > conf_thres # candidates xc = prediction[..., 4] > conf_thres # candidates
# Settings # Settings
...@@ -277,7 +277,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None, ...@@ -277,7 +277,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None,
merge = False # use merge-NMS merge = False # use merge-NMS
t = time.time() t = time.time()
output = [torch.zeros(0, 6)] * prediction.shape[0] output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
for xi, x in enumerate(prediction): # image index, image inference for xi, x in enumerate(prediction): # image index, image inference
# Apply constraints # Apply constraints
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论