def get_pred(temp, nms_th, nms_iou, oh, ow):
'''
temp:
cls_i_preds: LongTensor(batch_num, topk)
cls_p_preds: FloatTensor(batch_num, topk)
reg_preds: FloatTensor(batch_num, topk, 4)
Return:
cls_i_preds: (LongTensor(s1), LongTensor(s2), ...)
cls_p_preds: (FloatTensor(s1), FloatTensor(s2), ...)
reg_preds: (FloatTensor(s1,4), FloatTensor(s2,4), ...)
'''
cls_i_preds, cls_p_preds, reg_preds = temp
_cls_i_preds, _cls_p_preds, _reg_preds = [], [], []
for b in range(cls_i_preds.shape[0]):
cls_i_preds_b = cls_i_preds[b]
cls_p_preds_b = cls_p_preds[b]
reg_preds_b = reg_preds[b]
mask = cls_p_preds_b > nms_th
cls_i_preds_b = cls_i_preds_b[mask]
cls_p_preds_b = cls_p_preds_b[mask]
reg_preds_b = reg_preds_b[mask]
keep = box_nms(reg_preds_b, cls_p_preds_b, nms_iou)
cls_i_preds_b = cls_i_preds_b[keep]
cls_p_preds_b = cls_p_preds_b[keep]
reg_preds_b = reg_preds_b[keep]
reg_preds_b[:, :2] = reg_preds_b[:, :2].clamp_(min=0)
reg_preds_b[:, 2] = reg_preds_b[:, 2].clamp_(max=float(oh[b]) - 1)
reg_preds_b[:, 3] = reg_preds_b[:, 3].clamp_(max=float(ow[b]) - 1)
_cls_i_preds.append(cls_i_preds_b)
_cls_p_preds.append(cls_p_preds_b)
_reg_preds.append(reg_preds_b)
return _cls_i_preds, _cls_p_preds, _reg_preds
def get_pred(temp, nms_th, nms_iou):
'''
temp:
cls_i_preds: LongTensor(batch_num, topk)
cls_p_preds: FloatTensor(batch_num, topk)
reg_preds: FloatTensor(batch_num, topk, 4)
Return:
cls_i_preds: (LongTensor(s1), LongTensor(s2), ...)
cls_p_preds: (FloatTensor(s1), FloatTensor(s2), ...)
reg_preds: (FloatTensor(s1,4), FloatTensor(s2,4), ...)
'''
cls_i_preds, cls_p_preds, reg_preds = temp
_cls_i_preds, _cls_p_preds, _reg_preds = [], [], []
for b in range(cls_i_preds.shape[0]):
cls_i_preds_b = cls_i_preds[b]
cls_p_preds_b = cls_p_preds[b]
reg_preds_b = reg_preds[b]
mask = cls_p_preds_b > nms_th
cls_i_preds_b = cls_i_preds_b[mask]
cls_p_preds_b = cls_p_preds_b[mask]
reg_preds_b = reg_preds_b[mask]
keep = box_nms(reg_preds_b, cls_p_preds_b, nms_iou)
_cls_i_preds.append(cls_i_preds_b[keep])
_cls_p_preds.append(cls_p_preds_b[keep])
_reg_preds.append(reg_preds_b[keep])
return _cls_i_preds, _cls_p_preds, _reg_preds
def get_pred(temp, nms_th, nms_iou):
'''
temp:
pred_cls_i: L(b, topk)
pred_cls_p: F(b, topk)
pred_reg: F(b, topk, 4)
Return:
cls_i_preds: (L(s1), L(s2), ...)
cls_p_preds: (F(s1), F(s2), ...)
reg_preds: (F(s1,4), F(s2,4), ...)
'''
pred_cls_i, pred_cls_p, pred_reg = temp
_pred_cls_i, _pred_cls_p, _pred_reg = [], [], []
for b in range(pred_cls_i.shape[0]):
m = pred_cls_p[b] > nms_th
pred_cls_i_b = pred_cls_i[b][m]
pred_cls_p_b = pred_cls_p[b][m]
pred_reg_b = pred_reg[b][m]
keep = box_nms(pred_reg_b, pred_cls_p_b, nms_iou)
_pred_cls_i.append(pred_cls_i_b[keep])
_pred_cls_p.append(pred_cls_p_b[keep])
_pred_reg.append(pred_reg_b[keep])
return _pred_cls_i, _pred_cls_p, _pred_reg