import numpy as np
return_iou=True
iou = np.array([0.41, 0.2, 0.7, 0.6])
encoded_box = np.zeros((iou.shape[0], 4 + return_iou))
print(encoded_box)
print(iou.argmax())
assign_mask = iou > 0.5
print('-----------')
print(iou[assign_mask])
print(assign_mask)
if not assign_mask.any():
assign_mask[iou.argmax()] = True
if return_iou:
encoded_box[:, -1][assign_mask] = iou[assign_mask]
print(encoded_box)
priors = np.array([[0.2, 0.3, 0.5, 0.6], [0.3, 0.4, 0.5, 0.6], [0.3, 0.4, 0.4, 0.7], [0.1, 0.3, 0.2, 0.8]])
print(encoded_box[:, :2][assign_mask])
priors[:,:2]= np.array([0.2,0.2])+priors[:,:2]
print(priors[:,:2][assign_mask])
上述代码的输出:
[[0. 0. 0. 0. 0.]
[0. 0. 0. 0. 0.]
[0. 0. 0. 0. 0.]
[0. 0. 0. 0. 0.]]
2
-----------
[0.7 0.6]
[False False True True]
[[0. 0. 0. 0. 0. ]
[0. 0. 0. 0. 0. ]
[0. 0. 0. 0. 0.7]
[0. 0. 0. 0. 0.6]]
[[0. 0.]
[0. 0.]]
[[0.5 0.6]
[0.3 0.5]]