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
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
def run_on_image(img_url, net):
img = get_image(img_url)
input_blobs = _prepare_blobs(img, cfg.PIXEL_MEANS, cfg.TEST.SCALE,
cfg.TEST.MAX_SIZE)
for k, v in input_blobs.items():
workspace.FeedBlob(core.ScopedName(k), v,
mutils.get_device_option_cpu())
workspace.RunNetOnce(net)
blob_names = workspace.Blobs()
for x in blob_names:
print(x)
#goods = [x for x in blob_names if x.split('_')[-1] == 'nms']
#print(goods)
scores = workspace.FetchBlob('score_nms')
classids = workspace.FetchBlob('class_nms')
boxes = workspace.FetchBlob('bbox_nms')
cls_prob = workspace.FetchBlob('cls_prob')
bbox_pred = workspace.FetchBlob('bbox_pred')
print(scores)
print(classids)
print(boxes)
print("cls_prob: shape {}".format(cls_prob.shape))
print(cls_prob)
print("bbox_pred: shape {}".format(bbox_pred.shape))
print(bbox_pred)
#except Exception as e:
# print('Model failed to run')
# 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))
#print(boxes)
boxes, _, _, classids = _sort_results(boxes, None, None, classids)
check_blobs = [
"result_boxes",
"result_classids", # result
]
workspace.FeedBlob('result_boxes', boxes)
workspace.FeedBlob('result_classids', classids)
ret = _get_result_blobs(check_blobs)
return ret
def convert_model_gpu(args, net, init_net):
assert args.device == 'gpu'
ret_net = copy.deepcopy(net)
ret_init_net = copy.deepcopy(init_net)
cdo_cuda = mutils.get_device_option_cuda()
cdo_cpu = mutils.get_device_option_cpu()
CPU_OPS = [
["CollectAndDistributeFpnRpnProposals", None],
["GenerateProposals", None],
["BBoxTransform", None],
["BoxWithNMSLimit", None],
]
CPU_BLOBS = ["im_info", "anchor"]
@op_filter()
def convert_op_gpu(op):
for x in CPU_OPS:
if mutils.filter_op(op, type=x[0], inputs=x[1]):
return None
op.device_option.CopyFrom(cdo_cuda)
return [op]
@op_filter()
def convert_init_op_gpu(op):
if op.output[0] in CPU_BLOBS:
op.device_option.CopyFrom(cdo_cpu)
else:
op.device_option.CopyFrom(cdo_cuda)
return [op]
convert_op_in_proto(ret_init_net.Proto(), convert_init_op_gpu)
convert_op_in_proto(ret_net.Proto(), convert_op_gpu)
ret = core.InjectDeviceCopiesAmongNets([ret_init_net, ret_net])
return [ret[0][1], ret[0][0]]
def convert_model_gpu(args, net, init_net):
assert args.device == 'gpu'
ret_net = copy.deepcopy(net)
ret_init_net = copy.deepcopy(init_net)
cdo_cuda = mutils.get_device_option_cuda()
cdo_cpu = mutils.get_device_option_cpu()
CPU_OPS = [
["CollectAndDistributeFpnRpnProposals", None],
["GenerateProposals", None],
["BBoxTransform", None],
["BoxWithNMSLimit", None],
]
CPU_BLOBS = ["im_info", "anchor"]
@op_filter()
def convert_op_gpu(op):
for x in CPU_OPS:
if mutils.filter_op(op, type=x[0], inputs=x[1]):
return None
op.device_option.CopyFrom(cdo_cuda)
return [op]
@op_filter()
def convert_init_op_gpu(op):
if op.output[0] in CPU_BLOBS:
op.device_option.CopyFrom(cdo_cpu)
else:
op.device_option.CopyFrom(cdo_cuda)
return [op]
convert_op_in_proto(ret_init_net.Proto(), convert_init_op_gpu)
convert_op_in_proto(ret_net.Proto(), convert_op_gpu)
ret = core.InjectDeviceCopiesAmongNets([ret_init_net, ret_net])
return [ret[0][1], ret[0][0]]
