image_vis.py 文件源码

def draw_anchor_batch(pred_dict, image):
    """
    Draw anchors used in the batch for RPN.
    """
    anchors = pred_dict['all_anchors']
    targets = pred_dict['rpn_prediction']['rpn_cls_target']

    in_batch_idx = targets >= 0
    anchors = anchors[in_batch_idx]
    targets = targets[in_batch_idx]

    image_pil, draw = get_image_draw(image)

    for anchor, target in zip(anchors, targets):
        if target == 1:
            draw.rectangle(
                list(anchor), fill=(20, 200, 10, 15),
                outline=(20, 200, 10, 30))
        else:
            draw.rectangle(
                list(anchor), fill=(200, 10, 170, 10),
                outline=(200, 10, 170, 30))

    return image_pil