Detectron .pkl模型转成Caffe2 .pb模型

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)