image_vis.py 文件源码

def recalculate_objects(pred_dict, image):
    proposals = pred_dict['rpn_prediction']['proposals']
    proposals_prob = pred_dict['classification_prediction']['rcnn']['cls_prob']
    proposals_target = proposals_prob.argmax(axis=1) - 1
    bbox_offsets = pred_dict[
        'classification_prediction']['rcnn']['bbox_offsets']

    bbox_offsets = bbox_offsets[proposals_target >= 0]
    proposals = proposals[proposals_target >= 0]
    proposals_target = proposals_target[proposals_target >= 0]

    bbox_offsets_idx_pairs = np.stack(
        np.array([
            proposals_target * 4, proposals_target * 4 + 1,
            proposals_target * 4 + 2, proposals_target * 4 + 3]), axis=1)
    bbox_offsets = np.take(bbox_offsets, bbox_offsets_idx_pairs.astype(np.int))

    bboxes = decode(proposals, bbox_offsets)

    return bboxes, proposals_target