check conf score negative

pull/139/head
Ninenox 1 year ago
parent 57985c8d32
commit 2fb1e5dd8c
  1. 3
      models/torch_utils.py

@ -64,6 +64,8 @@ def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]):
iou_thres: float = 0.65
num_dets, bboxes, scores, labels = data[0][0], data[1][0], data[2][
0], data[3][0]
# check score negative
scores[scores < 0] = 1+scores[scores < 0]
nums = num_dets.item()
if nums == 0:
return bboxes.new_zeros((0, 4)), scores.new_zeros(
@ -71,6 +73,7 @@ def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]):
bboxes = bboxes[:nums]
scores = scores[:nums]
labels = labels[:nums]
# add nms
idx = nms(bboxes, scores, iou_thres)
bboxes, scores, labels = bboxes[idx], scores[idx], labels[idx]
return bboxes, scores, labels

Loading…
Cancel
Save