sample.py 文件源码

def sample_rpn_outputs_wrt_gt_boxes(boxes, scores, gt_boxes, is_training=False, only_positive=False):
    """sample boxes for refined output"""
    boxes, scores, batch_inds = sample_rpn_outputs(boxes, scores, is_training, only_positive)

    if gt_boxes.size > 0:
        overlaps = cython_bbox.bbox_overlaps(
                np.ascontiguousarray(boxes[:, 0:4], dtype=np.float),
                np.ascontiguousarray(gt_boxes[:, 0:4], dtype=np.float))
        gt_assignment = overlaps.argmax(axis=1) # B
        max_overlaps = overlaps[np.arange(boxes.shape[0]), gt_assignment] # B
        fg_inds = np.where(max_overlaps >= cfg.FLAGS.fg_threshold)[0]

        mask_fg_inds = np.where(max_overlaps >= cfg.FLAGS.mask_threshold)[0]
        if mask_fg_inds.size > cfg.FLAGS.masks_per_image:
            mask_fg_inds = np.random.choice(mask_fg_inds, size=cfg.FLAGS.masks_per_image, replace=False)

        if True:
            gt_argmax_overlaps = overlaps.argmax(axis=0) # G
            fg_inds = np.union1d(gt_argmax_overlaps, fg_inds)

    fg_rois = int(min(fg_inds.size, cfg.FLAGS.rois_per_image * cfg.FLAGS.fg_roi_fraction))
        if fg_inds.size > 0 and fg_rois < fg_inds.size:
           fg_inds = np.random.choice(fg_inds, size=fg_rois, replace=False)

    # TODO: sampling strategy
        bg_inds = np.where((max_overlaps < cfg.FLAGS.bg_threshold))[0]
        bg_rois = max(min(cfg.FLAGS.rois_per_image - fg_rois, fg_rois * 3), 64)
        if bg_inds.size > 0 and bg_rois < bg_inds.size:
           bg_inds = np.random.choice(bg_inds, size=bg_rois, replace=False)

    keep_inds = np.append(fg_inds, bg_inds)
    else:
        bg_inds = np.arange(boxes.shape[0])
        bg_rois = min(int(cfg.FLAGS.rois_per_image * (1-cfg.FLAGS.fg_roi_fraction)), 64)
        if bg_rois < bg_inds.size:
            bg_inds = np.random.choice(bg_inds, size=bg_rois, replace=False)

        keep_inds = bg_inds
        mask_fg_inds = np.arange(0)

    return boxes[keep_inds, :], scores[keep_inds], batch_inds[keep_inds],\
           boxes[mask_fg_inds, :], scores[mask_fg_inds], batch_inds[mask_fg_inds]