detectron目标检测.pkl模型转.pb模型已经将生成的pkl模型转换成为了pb模型，那么如何使用呢，代码如下：
import cv2
import numpy as np
from caffe2.python import core, workspace
from caffe2.proto import caffe2_pb2
import time
from caffe2.python import dyndep
#需要detectron中一个库的支持，因此必须导入，该库位于caffe2的路径中
detectron_ops_lib = '/home/yantianwang/clone/pytorch/build/lib/libcaffe2_detectron_ops_gpu.so'
dyndep.InitOpsLibrary(detectron_ops_lib)
def get_device_option_cpu():
device_option = core.DeviceOption(caffe2_pb2.CPU)
return device_option
def get_device_option_cuda(gpu_id=0):
device_option = caffe2_pb2.DeviceOption()
device_option.device_type = caffe2_pb2.CUDA
device_option.cuda_gpu_id = gpu_id
return device_option
def _sort_results(boxes, segms, keypoints, classes):
indices = np.argsort(boxes[:, -1])[::-1]
if boxes is not None:
boxes = boxes[indices, :]
if segms is not None:
segms = [segms[x] for x in indices]
if keypoints is not None:
keypoints = [keypoints[x] for x in indices]
if classes is not None:
if isinstance(classes, list):
classes = [classes[x] for x in indices]
else:
classes = classes[indices]
return boxes, segms, keypoints, classes
def im_list_to_blob(ims):
"""Convert a list of images into a network input. Assumes images were
prepared using prep_im_for_blob or equivalent: i.e.
- BGR channel order
- pixel means subtracted
- resized to the desired input size
- float32 numpy ndarray format
Output is a 4D HCHW tensor of the images concatenated along axis 0 with
shape.
"""
if not isinstance(ims, list):
ims = [ims]
max_shape = np.array([im.shape for im in ims]).max(axis=0)
# Pad the image so they can be divisible by a stride
if FPN_ON:
stride = float(FPN_COARSEST_STRIDE)
max_shape[0] = int(np.ceil(max_shape[0] / stride) * stride)
max_shape[1] = int(np.ceil(max_shape[1] / stride) * stride)
num_images = len(ims)
blob = np.zeros(
(num_images, max_shape[0], max_shape[1], 3), dtype=np.float32
)
for i in range(num_images):
im = ims[i]
blob[i, 0:im.shape[0], 0:im.shape[1], :] = im
# Move channels (axis 3) to axis 1
# Axis order will become: (batch elem, channel, height, width)
channel_swap = (0, 3, 1, 2)
blob = blob.transpose(channel_swap)
return blob
def _prepare_blobs(
im,
pixel_means,
target_size,
max_size,
):
''' Reference: blob.prep_im_for_blob() '''
im = im.astype(np.float32, copy=False)
im -= pixel_means
im_shape = im.shape
im_size_min = np.min(im_shape[0:2])
im_size_max = np.max(im_shape[0:2])
im_scale = float(target_size) / float(im_size_min)
if np.round(im_scale * im_size_max) > max_size:
im_scale = float(max_size) / float(im_size_max)
im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale,
interpolation=cv2.INTER_LINEAR)
# Reuse code in blob_utils and fit FPN
blob = im_list_to_blob([im])
blobs = {}
blobs['data'] = blob
blobs['im_info'] = np.array(
[[blob.shape[2], blob.shape[3], im_scale]],
dtype=np.float32
)
return blobs
def create_input_blobs_for_net(net_def):
for op in net_def.op:
for blob_in in op.input:
if not workspace.HasBlob(blob_in):
workspace.CreateBlob(blob_in)
print("Current blobs in the workspace: {}".format(workspace.Blobs()))
def run_model_pb(net, init_net, im):
workspace.ResetWorkspace()
workspace.RunNetOnce(init_net)
create_input_blobs_for_net(net)
workspace.CreateNet(net)
PIXEL_MEANS = np.array([[[102.9801, 115.9465, 122.7717]]])
# input_blobs, _ = core_test._get_blobs(im, None)
input_blobs = _prepare_blobs(
im,
PIXEL_MEANS,
TEST_SCALE, TEST_MAX_SIZE
)
gpu_blobs = []
gpu_blobs = ['data']
#GPU
for k, v in input_blobs.items():
workspace.FeedBlob(
core.ScopedName(k),
v,
get_device_option_cuda() if k in gpu_blobs else
get_device_option_cpu()
)
#CPU
#for k, v in input_blobs.items():
# workspace.FeedBlob(
# core.ScopedName(k),
# v,
# device_option = core.DeviceOption(caffe2_pb2.CPU)
#)
try:
workspace.RunNet(net.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)
return boxes,classids
if __name__ == '__main__':
FPN_ON = True
FPN_COARSEST_STRIDE = 32
TEST_SCALE, TEST_MAX_SIZE = 1024,1024
test_img_file = '456.png'
print('Loading test file {}...'.format(test_img_file))
test_img = cv2.imread(test_img_file)
assert test_img is not None
predict_net = caffe2_pb2.NetDef()
init_net = caffe2_pb2.NetDef()
with open("model.pb") as f:
predict_net.ParseFromString(f.read())
with open("model_init.pb") as f:
init_net.ParseFromString(f.read())
#print(predict_net)
print(predict_net.name)
print(init_net.name)
print('1')
start = time.time()
boxes,classids = run_model_pb(predict_net, init_net, test_img)
end = time.time()
print(end-start)
print(boxes)这里打印的是检测出来的boxes