tadax / m2det

M2Det implementation using TensorFlow
MIT License
63 stars 26 forks source link

about assign_boxes(boxes, priors, num_classes, threshold=0.5): #12

Open bjmajic opened 5 years ago

bjmajic commented 5 years ago

def assign_boxes(boxes, priors, num_classes, threshold=0.5): num_classes += 1 # add background class assignment = np.zeros((len(priors), 4 + num_classes + 1)) assignment[:, 4] = 1.0 # background encoded_boxes = np.apply_along_axis(encode_box, 1, boxes[:, :4], priors, threshold) encoded_boxes = encoded_boxes.reshape(-1, len(priors), 5) best_iou = encoded_boxes[:, :, -1].max(axis=0) best_iou_idx = encoded_boxes[:, :, -1].argmax(axis=0)

I think the axis should be 1 , not 0. it that right? why?

best_iou_mask = best_iou > 0 # judge by iou between prior and bbox
best_iou_idx = best_iou_idx[best_iou_mask]
assign_num = len(best_iou_idx)
encoded_boxes = encoded_boxes[:, best_iou_mask, :]
assignment[:, :4][best_iou_mask] = encoded_boxes[best_iou_idx, np.arange(assign_num), :4]
assignment[:, 4][best_iou_mask] = 0 # background
assignment[:, 5:-1][best_iou_mask] = boxes[best_iou_idx, 4:]
assignment[:, -1][best_iou_mask] = 1 # objectness

return assignment