Detectron训练出来的目标检测模型后缀为.pkl，在使用的时候必须要有图（graph）以及Detectron代码的支持，转为caffe2标准的pb模型后，就可以脱离detectron的代码单独运行。
转换用到的 tools/convert_pkl_to_pb.py
1.png
2.png
3.png
生成三个文件
并在对应的yaml文件中加入
4.png
测试代码.pb模型代码,根据自己情况进行修改
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName :pb_infer.py
# @Time :2020/6/2 15:42
# @Author :wanlin
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
import os
# 需要detectron中一个库的支持，因此必须导入，该库位于caffe2的路径中
detectron_ops_lib = '/home/wl/anaconda3/envs/wl_torch14/lib/python3.6/site-packages/torch/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.device_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
FPN_COARSEST_STRIDE = 32
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)
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,
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]]])
# PIXEL_MEANS = np.array([[[0, 0, 0]]])
# input_blobs, _ = core_test._get_blobs(im, None)
input_blobs = _prepare_blobs(
im,
TEST_SCALE, TEST_MAX_SIZE
)
# gpu_blobs = []
gpu_blobs = ['data']
# GPU
t1 = time.time()
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)
t2 = time.time()
print('spend time: ', t2-t1)
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__':
model_path = './pb_file'
FPN_ON = True
# FPN_COARSEST_STRIDE = 32
TEST_SCALE, TEST_MAX_SIZE = 800, 1024
test_img_file = './test_image.jpg'
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(os.path.join(model_path,"model.pb"),'rb') as f:
predict_net.ParseFromString(f.read())
with open(os.path.join(model_path,"model_init.pb"), 'rb') as f:
init_net.ParseFromString(f.read())
label_color = {1:(255,0,0),2:(0,0,255),3:(0,255,0)}
start = time.time()
boxes, classids = run_model_pb(predict_net, init_net, test_img)
end = time.time()
print(end - start)
for box, classid in zip(boxes, classids):
if box[4] >= 0.7:
cv2.rectangle(test_img, (box[0],box[1]),(box[2],box[3]),label_color[int(classid)],3)
cv2.imwrite('./test_image_result.jpg', test_img)