onnx转rknn推理中的filter_boxes函数问题

[复制链接] · 发表于 2023-5-17 09:26:28

def filter_boxes(boxes, box_confidences, box_class_probs):
"""Filter boxes with box threshold. It's a bit different with origin yolov5 post process!

      # Arguments
         boxes: ndarray, boxes of objects.
         box_confidences: ndarray, confidences of objects.
         box_class_probs: ndarray, class_probs of objects.

      # Returns
         boxes: ndarray, filtered boxes.
         classes: ndarray, classes for boxes.
         scores: ndarray, scores for boxes.
"""
boxes = boxes.reshape(-1, 4)
box_confidences = box_confidences.reshape(-1)
box_class_probs = box_class_probs.reshape(-1, box_class_probs.shape[-1])

_box_pos = np.where(box_confidences >= OBJ_THRESH)
boxes = boxes[_box_pos]
box_confidences = box_confidences[_box_pos]
box_class_probs = box_class_probs[_box_pos]

class_max_score = np.max(box_class_probs, axis=-1)
classes = np.argmax(box_class_probs, axis=-1)
_class_pos = np.where(class_max_score >= OBJ_THRESH)

boxes = boxes[_class_pos]
classes = classes[_class_pos]
scores = (class_max_score* box_confidences)[_class_pos]

return boxes, classes, scores

def filter_boxes(boxes, box_confidences, box_class_probs):
"""Filter boxes with box threshold. It's a bit different with origin yolov5 post process!

# Arguments
      boxes: ndarray, boxes of objects.
      box_confidences: ndarray, confidences of objects.
      box_class_probs: ndarray, class_probs of objects.

# Returns
      boxes: ndarray, filtered boxes.
      classes: ndarray, classes for boxes.
      scores: ndarray, scores for boxes.
"""

box_classes = np.argmax(box_class_probs, axis=-1)
box_class_scores = np.max(box_class_probs, axis=-1)
pos = np.where(box_confidences[..., 0] >= BOX_THESH)

boxes = boxes[pos]
classes = box_classes[pos]
scores = box_class_scores[pos]
return boxes, classes, scores

第一个函数代码是rknn-toolkit2\examples\onnx\yolov5\test.py里的，第二个函数代码是网上部分人使用的，请问第二个函数和第一个有什么区别，为什么采用第二个函数输出框的置信度高一些。