예제 #1
0
def run_model_pb(args, net, init_net, im, check_blobs):
    workspace.ResetWorkspace()
    workspace.RunNetOnce(init_net)
    mutils.create_input_blobs_for_net(net.Proto())
    workspace.CreateNet(net)

    # input_blobs, _ = core_test._get_blobs(im, None)
    input_blobs = _prepare_blobs(
        im,
        cfg.PIXEL_MEANS,
        cfg.TEST.SCALE, cfg.TEST.MAX_SIZE
    )
    gpu_blobs = []
    if args.device == 'gpu':
        gpu_blobs = ['data']
    for k, v in input_blobs.items():
        workspace.FeedBlob(
            core.ScopedName(k),
            v,
            mutils.get_device_option_cuda()
            if k in gpu_blobs
            else mutils.get_device_option_cpu()
        )

    try:
        workspace.RunNet(net)
        scores = workspace.FetchBlob('score_nms')
        classids = workspace.FetchBlob('class_nms')
        boxes = workspace.FetchBlob('bbox_nms')
    except Exception as e:
        print('Running pb model failed.\n{}'.format(e))
        # may not detect anything at all
        R = 0
        scores = np.zeros((R,), dtype=np.float32)
        boxes = np.zeros((R, 4), dtype=np.float32)
        classids = np.zeros((R,), dtype=np.float32)

    boxes = np.column_stack((boxes, scores))

    # sort the results based on score for comparision
    boxes, _, _, classids = _sort_results(
        boxes, None, None, classids)

    # write final result back to workspace
    workspace.FeedBlob('result_boxes', boxes)
    workspace.FeedBlob('result_classids', classids)

    ret = _get_result_blobs(check_blobs)

    return ret
예제 #2
0
def run_model_pb(args, net, init_net, im, check_blobs):
    workspace.ResetWorkspace()
    workspace.RunNetOnce(init_net)
    mutils.create_input_blobs_for_net(net.Proto())
    workspace.CreateNet(net)

    # input_blobs, _ = core_test._get_blobs(im, None)
    input_blobs = _prepare_blobs(
        im,
        cfg.PIXEL_MEANS,
        cfg.TEST.SCALE, cfg.TEST.MAX_SIZE
    )
    gpu_blobs = []
    if args.device == 'gpu':
        gpu_blobs = ['data']
    for k, v in input_blobs.items():
        workspace.FeedBlob(
            core.ScopedName(k),
            v,
            mutils.get_device_option_cuda() if k in gpu_blobs else
            mutils.get_device_option_cpu()
        )

    try:
        workspace.RunNet(net.Proto().name)
        scores = workspace.FetchBlob('score_nms')
        classids = workspace.FetchBlob('class_nms')
        boxes = workspace.FetchBlob('bbox_nms')
    except Exception as e:
        print('Running pb model failed.\n{}'.format(e))
        # may not detect anything at all
        R = 0
        scores = np.zeros((R,), dtype=np.float32)
        boxes = np.zeros((R, 4), dtype=np.float32)
        classids = np.zeros((R,), dtype=np.float32)

    boxes = np.column_stack((boxes, scores))

    # sort the results based on score for comparision
    boxes, _, _, classids = _sort_results(
        boxes, None, None, classids)

    # write final result back to workspace
    workspace.FeedBlob('result_boxes', boxes)
    workspace.FeedBlob('result_classids', classids)

    ret = _get_result_blobs(check_blobs)

    return ret