Multi-Proposal Operator Test (cpu implementation) for MXNet

test_multi_proposal.py

import mxnet as mx
from mxnet import nd
import numpy as np

feature_stride = 16 
scales = (8, 16, 32)
ratios = (0.5, 1, 2)
rpn_pre_nms_top_n = 12000 
rpn_post_nms_top_n = 2000
threshold = 0.7
rpn_min_size = 16

batch_size = 20
feat_len = 14
H, W = feat_len, feat_len
num_anchors = len(scales) * len(ratios)
count_anchors = feat_len * feat_len * num_anchors

'''
    cls_prob: (batch_size, 2 * num_anchors, H, W)
    bbox_pred: (batch_size, 4 * num_anchors, H, W) 
    im_info: (batch_size, 3)
'''

cls_prob = nd.empty((batch_size, 2 * num_anchors, H, W), dtype = np.float32) 
bbox_pred = nd.empty((batch_size, 4 * num_anchors, H, W), dtype = np.float32) 
im_info = nd.empty((batch_size, 3), dtype = np.float32) 

cls_prob = nd.array(np.random.random(cls_prob.shape))
bbox_pred = nd.array(np.random.random(bbox_pred.shape))

for i in range(batch_size):
    im_size = np.random.randint(100, feat_len * feature_stride, size = (2,))
    im_scale = np.random.randint(70, 100) / 100.0
    im_info[i, :] = [im_size[0], im_size[1], im_scale]

def get_sub(arr, i):
    new_shape = list(arr.shape)
    new_shape[0] = 1
    res = arr[i].reshape(new_shape)
    return res

def test(rpn_pre_nms_top_n, rpn_post_nms_top_n):

    single_proposal = []
    single_score = []
    for i in range(batch_size):
        rois, score = mx.nd.contrib.Proposal(
                cls_score = get_sub(cls_prob, i),
                bbox_pred = get_sub(bbox_pred, i),
                im_info = get_sub(im_info, i),
                feature_stride = feature_stride,
                scales = scales,
                ratios = ratios,
                rpn_pre_nms_top_n = rpn_pre_nms_top_n,
                rpn_post_nms_top_n = rpn_post_nms_top_n,
                threshold = threshold,
                rpn_min_size = rpn_min_size, output_score = True)
        single_proposal.append(rois)
        single_score.append(score)

    multi_proposal, multi_score = mx.nd.contrib.MultiProposal(
            cls_score = cls_prob,
            bbox_pred = bbox_pred,
            im_info = im_info,
            feature_stride = feature_stride,
            scales = scales,
            ratios = ratios,
            rpn_pre_nms_top_n = rpn_pre_nms_top_n,
            rpn_post_nms_top_n = rpn_post_nms_top_n,
            threshold = threshold,
            rpn_min_size = rpn_min_size, output_score = True)

    single_proposal = nd.stack(*single_proposal).reshape(multi_proposal.shape)
    single_score = nd.stack(*single_score).reshape(multi_score.shape)

    single_proposal_np = single_proposal.asnumpy()
    multi_proposal_np = multi_proposal.asnumpy()

    single_score_np = single_score.asnumpy()
    multi_score_np = multi_score.asnumpy()
    print (multi_proposal_np.shape)

    # test rois x1,y1,x2,y2
    assert np.allclose(single_proposal_np[:, 1:], multi_proposal_np[:, 1:])
    # test rois batch_idx 
    for i in range(batch_size):
        start = i * rpn_post_nms_top_n
        end = start + rpn_post_nms_top_n
        assert (multi_proposal_np[start:end, 0] == i).all()

    # test score
    assert np.allclose(single_score_np, multi_score_np)

test(rpn_pre_nms_top_n, rpn_post_nms_top_n)
test(rpn_pre_nms_top_n, 1500)
test(1000, 500)
print ("test ok")