object detection with the onnx tensorrt backbone in python-yolov3_onnx

yolov3_to_onnx.py

from __future__ import print_function
from collections import OrderedDict
import hashlib
import os.path

import wget

import onnx
from onnx import helper
from onnx import TensorProto
import numpy as np

import sys

sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
from common import retry_call

class DarkNetParser(object):
    def __init__(self, supported_layers):
        self.layer_configs = OrderedDict()
        self.supported_layers = supported_layers
        self.layer_counter = 0
    def parse_cfg_file(self, cfg_file_path):
        with open(cfg_file_path) as cfg_file:
            remainder = cfg_file.read()
            while remainder is not None:
                layer_dict, layer_name, remainder = self._next_layer(remainder)
                if layer_dict is not None:
                    self.layer_configs[layer_name] = layer_dict
        return self.layer_configs
    def _next_layer(self, remainder):
        remainder = remainder.split('[', 1)
        if len(remainder) == 2:
            remainder = remainder[1]
        else:
            return None, None, None
        remainder = remainder.split(']', 1)
        if len(remainder) == 2:
            layer_type, remainder = remainder
        else:
            return None, None, None
        if remainder.replace(' ','')[0] == '#':
            remainder = remainder.split('\n', 1)[1]
        layer_param_block, remainder = remainder.split('\n\n',1)
        layer_param_lines = layer_param_block.split('\n')[1:]
        layer_name = str(self.layer_counter).zfill(3) + '_' + layer_type
        layer_dict = dict(type=layer_type)
        if layer_type in self.supported_layers:
            for param_line in layer_param_lines:
                if param_line[0] == '#':
                    continue
                param_type, param_value = self._parse_params(param_line)
                layer_dict[param_type] = param_value
        self.layer_counter += 1
        return layer_dict, layer_name, remainder
    def _parse_params(self, param_line):
        param_line = param_line.replace(' ','')
        param_type, param_value_raw = param_line.split('=')
        param_value = None
        if param_type == 'layers':
            layer_indexes = list()
            for index in param_value_raw.split(','):
                layer_indexes.append(int(index))
            param_value = layer_indexes
        elif isinstance(param_value_raw, str) and not param_value_raw.isalpha():
            condition_param_value_positive = param_value_raw.isdigit()
            condition_param_value_negative = param_value_raw[0] == '-' and param_value_raw[1:].isdigit()
            if condition_param_value_positive or condition_param_value_negative:
                param_value = int(param_value_raw)
            else:
                param_value = float(param_value_raw)
        else:
            param_value = str(param_value_raw)
        return param_type, param_value
class MajorNodeSpecs(object):
    def __init__(self, name, channels):
        self.name = name
        self.channels = channels
        self.created_onnx_node = False
        if name is not None and isinstance(channels, int) and channels>0:
            self.created_onnx_node = True
class ConvParams(object):
    def __init__(self, node_name, batch_normalize, conv_weight_dims):
        self.node_name = node_name
        self.batch_normalize = batch_normalize
        assert len(conv_weight_dims) == 4
        self.conv_weight_dims = conv_weight_dims
    def generate_param_name(self, param_category, suffix):
        assert suffix
        assert param_category in ['bn', 'conv']
        assert (suffix in ['scale', 'mean', 'var', 'weights', 'bias'])
        if param_category == 'bn':
            assert self.batch_normalize
            assert  suffix in ['scale', 'bias', 'mean', 'var']
        elif param_category == 'conv':
            assert suffix in ['weights', 'bias']
            if suffix == 'bias':
                assert not self.batch_normalize
        param_name = self.node_name + '_' + param_category + '_' + suffix
        return param_name
class ResizeParams(object):
    def __init__(self, node_name, value):
        self.node_name = node_name
        self.value = value
    def generate_param_name(self):
        param_name = self.node_name + '_' + "scale"
        return param_name
    def generate_roi_name(self):
        param_name = self.node_name + '_' + "roi"
        return param_name
class WeightLoader(object):
    def __init__(self, weights_file_path):
        self.weights_file = self._open_weights_file(weights_file_path)
    def load_resize_scales(self, resize_params):
        initializer = list()
        inputs = list()
        name = resize_params.generate_param_name()
        shape = resize_params.value.shape
        data = resize_params.value
        scale_init = helper.make_tensor(
            name, TensorProto.FLOAT, shape, data
        )
        scale_input = helper.make_tensor_value_info(
            name, TensorProto.FLOAT, shape
        )
        initializer.append(scale_init)
        inputs.append(scale_input)
        rank =4
        roi_name = resize_params.generate_roi_name()
        roi_input = helper.make_tensor_value_info(roi_name, TensorProto.FLOAT, [rank])
        roi_init = helper.make_tensor(roi_name, TensorProto.FLOAT, [rank], [0,0,0,0])
        initializer.append(roi_init)
        inputs.append(roi_input)
        return initializer, inputs
    def load_conv_weights(self, conv_params):
        initializer = list()
        inputs = list()
        if conv_params.batch_normalize:
            bias_init, bias_input = self._create_param_tensors(conv_params, 'bn', 'bias')
            bn_scale_init, bn_scale_input = self._create_param_tensors(conv_params, 'bn', 'scale')
            bn_mean_init, bn_mean_input = self._create_param_tensors(conv_params,'bn','mean')
            bn_var_init, bn_var_input = self._create_param_tensors(conv_params, 'bn', 'var')
            initializer.extend([bn_scale_init, bias_init, bn_mean_init, bn_var_init])
            inputs.extend([bn_scale_input, bias_input, bn_mean_input, bn_var_input])
        else:
            bias_init, bias_input = self._create_param_tensors(conv_params, 'conv', 'bias')
            initializer.append(bias_init)
            inputs.append(bias_input)
        conv_init, conv_input = self._create_param_tensors(conv_params, 'conv', 'weights')
        initializer.append(conv_init)
        inputs.append(conv_input)
        return initializer, inputs
    def _open_weights_file(self, weights_file_path):
        weights_file = open(weights_file_path, 'rb')
        length_header = 5
        np.ndarray(shape=(length_header,),dtype='int32',buffer=weights_file.read(
            length_header*4
        ))
        return weights_file
    def _create_param_tensors(self,conv_params, param_category, suffix):
        param_name, param_data, param_data_shape = self._load_one_param_type(
            conv_params, param_category, suffix
        )
        initializer_tensor = helper.make_tensor(param_name, TensorProto.FLOAT, param_data_shape, param_data)
        input_tensor = helper.make_tensor_value_info(param_name, TensorProto.FLOAT, param_data_shape)
        return initializer_tensor, input_tensor
    def _load_one_param_type(self, conv_params, param_category, suffix):
        param_name = conv_params.generate_param_name(param_category, suffix)
        channels_out, channels_in, filter_h, filter_w = conv_params.conv_weight_dims
        if param_category == 'bn':
            param_shape = [channels_out]
        elif param_category == 'conv':
            if suffix == 'weights':
                param_shape = [channels_out, channels_in, filter_h, filter_w]
            elif suffix == 'bias':
                param_shape = [channels_out]
        param_size = np.product(np.array(param_shape))
        param_data = np.ndarray(
            shape = param_shape,
            dtype = 'float32',
            buffer = self.weights_file.read(param_size * 4)
        )
        param_data = param_data.flatten().astype(float)
        return param_name, param_data, param_shape
class GraphBuilderONNX(object):
    def __init__(self, output_tensors):
        self.output_tensors = output_tensors
        self._nodes = list()
        self.graph_def = None
        self.input_tensor = None
        self.epsilon_bn = 1e-5
        self.momentum_bn = 0.99
        self.alpha_lrelu = 0.1
        self.param_dict = OrderedDict()
        self.major_node_specs = list()
        self.batch_size = 1
    def build_onnx_graph(self,
                         layer_configs,
                         weights_file_path,
                         verbose=True):
        for layer_name in layer_configs.keys():
            layer_dict = layer_configs[layer_name]
            major_node_specs = self._make_onnx_node(layer_name, layer_dict)
            if major_node_specs.name is not None:
                self.major_node_specs.append(major_node_specs)
        outputs = list()
        for tensor_name in self.output_tensors.keys():
            output_dims = [self.batch_size, ]+ self.output_tensors[tensor_name]
            output_tensor = helper.make_tensor_value_info(
                tensor_name, TensorProto.FLOAT, output_dims
            )
            outputs.append(output_tensor)
        inputs = [self.input_tensor]
        weight_loader = WeightLoader(weights_file_path)
        initializer = list()
        for layer_name in self.param_dict.keys():
            _, layer_type = layer_name.split('_', 1)
            params = self.param_dict[layer_name]
            if layer_type == 'convolutional':
                initializer_layer, inputs_layer = weight_loader.load_conv_weights(params)
                initializer.extend(initializer_layer)
                inputs.extend(inputs_layer)
            elif layer_type == 'upsample':
                initializer_layer, inputs_layer = weight_loader.load_resize_scales(params)
                initializer.extend(initializer_layer)
                inputs.extend(inputs_layer)
        del weight_loader
        self.graph_def = helper.make_graph(
            nodes = self._nodes,
            name = 'YOLOv3-608',
            inputs=inputs,
            outputs=outputs,
            initializer=initializer
        )
        if verbose:
            print(helper.printable_graph(self.graph_def))
        model_def = helper.make_model(self.graph_def, producer_name = 'NVIDIA TensorRT sample')
        return model_def
    def _make_onnx_node(self, layer_name, layer_dict):
        layer_type = layer_dict['type']
        if self.input_tensor is None:
            if layer_type == 'net':
                major_node_output_name, major_node_output_channels = self._make_input_tensor(layer_name, layer_dict)
                major_node_specs = MajorNodeSpecs(major_node_output_name, major_node_output_channels)
            else:
                raise ValueError('The first node has to be of type "net" .')
        else:
            node_creators = dict()
            node_creators['convolutional'] = self._make_conv_node
            node_creators['shortcut'] = self._make_shortcut_node
            node_creators['route'] = self._make_route_node
            node_creators['upsample'] = self._make_resize_node

            if layer_type in node_creators.keys():
                major_node_output_name, major_node_output_channels = node_creators[layer_type](layer_name, layer_dict)
                major_node_specs = MajorNodeSpecs(major_node_output_name, major_node_output_channels)
            else:
                print(
                    'Layer of type %s not supported, skipping ONNX node generation.' %
                    layer_type)
                major_node_specs = MajorNodeSpecs(layer_name, None)
        return major_node_specs
    def _make_input_tensor(self, layer_name, layer_dict):
        batch_size = layer_dict['batch']
        channels = layer_dict['channels']
        height = layer_dict['height']
        width = layer_dict['width']
        self.batch_size = batch_size
        input_tensor = helper.make_tensor_value_info(
            str(layer_name), TensorProto.FLOAT, [
                batch_size, channels, height, width
            ]
        )
        self.input_tensor = input_tensor
        return layer_name, channels
    def _get_previous_node_specs(self, target_index = -1):
        previous_node = None
        for node in self.major_node_specs[target_index::-1]:
            if node.created_onnx_node:
                previous_node = node
                break
        assert previous_node is not None
        return previous_node
    def _make_conv_node(self, layer_name, layer_dict):
        previous_node_specs = self._get_previous_node_specs()
        inputs = [previous_node_specs.name]
        previous_channels = previous_node_specs.channels
        kernel_size = layer_dict['size']
        stride = layer_dict['stride']
        filters = layer_dict['filters']
        batch_normalize = False
        if 'batch_normalize' in layer_dict.keys() and layer_dict['batch_normalize'] == 1:
            batch_normalize = True
        kernel_shape = [kernel_size, kernel_size]
        weights_shape = [filters, previous_channels] + kernel_shape
        conv_params = ConvParams(layer_name, batch_normalize, weights_shape)
        strides = [stride, stride]
        dilations = [1, 1]
        weights_name = conv_params.generate_param_name('conv', 'weights')
        inputs.append(weights_name)
        if not batch_normalize:
            bias_name = conv_params.generate_param_name('conv', 'bias')
            inputs.append(bias_name)
        conv_node = helper.make_node(
            'Conv',
            inputs= inputs,
            outputs=[layer_name],
            kernel_shape = kernel_shape,
            strides=strides,
            auto_pad='SAME_LOWER',
            dilations=dilations,
            name=layer_name
        )
        self._nodes.append(conv_node)
        inputs = [layer_name]
        layer_name_output = layer_name

        if batch_normalize:
            layer_name_bn = layer_name + '_bn'
            bn_param_suffixes = ['scale', 'bias', 'mean', 'var']
            for suffix in bn_param_suffixes:
                bn_param_name = conv_params.generate_param_name('bn', suffix)
                inputs.append(bn_param_name)
            batchnorm_node = helper.make_node(
                'BatchNormalization',
                inputs = inputs,
                outputs = [layer_name_bn],
                epsilon = self.epsilon_bn,
                momentum=self.momentum_bn,
                name = layer_name_bn
            )
            self._nodes.append(batchnorm_node)
            inputs = [layer_name_bn]
            layer_name_output = layer_name_bn
        if layer_dict['activation'] == 'leaky':
            layer_name_lrelu = layer_name + '_lrelu'
            lrelu_node = helper.make_node(
                'LeakyRelu',
                inputs=inputs,
                outputs=[layer_name_lrelu],
                name= layer_name_lrelu,
                alpha=self.alpha_lrelu
            )
            self._nodes.append(lrelu_node)
            inputs = [layer_name_lrelu]
            layer_name_output = layer_name_lrelu
        elif layer_dict['activation'] == 'linear':
            pass
        else:
            print('Activation not supported.')
        self.param_dict[layer_name] = conv_params
        return layer_name_output, filters
    def _make_shortcut_node(self, layer_name, layer_dict):
        shortcut_index = layer_dict['from']
        activation = layer_dict['activation']
        assert activation == 'linear'

        first_node_specs = self._get_previous_node_specs()
        second_node_specs = self._get_previous_node_specs(target_index = shortcut_index)
        assert first_node_specs.channels == second_node_specs.channels
        channels = first_node_specs.channels
        inputs = [first_node_specs.name, second_node_specs.name]
        shortcut_node = helper.make_node(
            'Add',
            inputs=inputs,
            outputs=[layer_name],
            name=layer_name,
        )
        self._nodes.append(shortcut_node)
        return layer_name, channels
    def _make_route_node(self, layer_name, layer_dict):
        route_node_indexes = layer_dict['layers']
        if len(route_node_indexes) == 1:
            split_index = route_node_indexes[0]
            assert split_index < 0
            split_index += 1
            self.major_node_specs = self.major_node_specs[:split_index]
            layer_name = None
            channels = None
        else:
            inputs = list()
            channels = 0
            for index in route_node_indexes:
                if index>0:
                    index += 1
                route_node_specs = self._get_previous_node_specs(target_index= index)
                inputs.append(route_node_specs.name)
                channels += route_node_specs.channels
            assert inputs
            assert channels > 0
            route_node = helper.make_node('Concat', axis =1, inputs = inputs, outputs=[layer_name], name = layer_name)
            self._nodes.append(route_node)
        return layer_name, channels
    def _make_resize_node(self, layer_name, layer_dict):
        resize_scale_factors = float(layer_dict['stride'])
        scales = np.array([1.0, 1.0, resize_scale_factors, resize_scale_factors]).astype(np.float32)
        previous_node_specs = self._get_previous_node_specs()
        inputs = [previous_node_specs.name]
        channels = previous_node_specs.channels
        assert channels>0
        resize_params = ResizeParams(layer_name, scales)
        roi_name = resize_params.generate_roi_name()
        inputs.append(roi_name)
        scales_name = resize_params.generate_param_name()
        inputs.append(scales_name)
        resize_node = helper.make_node('Resize', coordinate_transformation_mode = 'asymmetric', mode = 'nearest',
                                       nearest_mode='floor', inputs=inputs, outputs=[layer_name], name=layer_name)
        self._nodes.append(resize_node)
        self.param_dict[layer_name] = resize_params
        return layer_name, channels
def generate_md5_checksum(local_path):
    with open(local_path, 'rb') as local_file:
        data = local_file.read()
        return  hashlib.md5(data).hexdigest()
def download_file(local_path, link, checksum_reference=None):
    if not os.path.exists(local_path):
        print('Downloading from %s, this may take a while...' % link)
        retry_call(wget.download, args=[link, local_path], n_retries=3)
        print()
    if checksum_reference is not None:
        checksum = generate_md5_checksum(local_path)
        if checksum != checksum_reference:
            raise ValueError(
                'The MD5 checksum of local file %s differs from %s, please manually remove \
                 the file and try again.' %
                (local_path, checksum_reference))
    return local_path
def main():
    cfg_file_path = download_file('yolov3.cfg',
        'https://raw.githubusercontent.com/pjreddie/darknet/f86901f6177dfc6116360a13cc06ab680e0c86b0/cfg/yolov3.cfg',
        'b969a43a848bbf26901643b833cfb96c')
    supported_layers = ['net', 'convolutional', 'shortcut', 'route', 'upsample']
    parser = DarkNetParser(supported_layers)
    layer_configs = parser.parse_cfg_file(cfg_file_path)
    del parser
    output_tensor_dims = OrderedDict()
    output_tensor_dims['082_convolutional'] = [255, 19, 19]
    output_tensor_dims['094_convolutional'] = [255, 38, 38]
    output_tensor_dims['106_convolutional'] = [255, 76, 76]
    builder = GraphBuilderONNX(output_tensor_dims)
    weights_file_path = download_file(
        'yolov3.weights',
        'https://pjreddie.com/media/files/yolov3.weights',
        'c84e5b99d0e52cd466ae710cadf6d84c')
    yolov3_model_def = builder.build_onnx_graph(layer_configs=layer_configs, weights_file_path=weights_file_path, verbose=True)
    del builder
    onnx.checker.check_model(yolov3_model_def)
    output_file_path = 'yolov3.onnx'
    onnx.save(yolov3_model_def,output_file_path)
if __name__ == '__main__':
    main()



        

onnx_to_tensorrt.py

from __future__ import print_function
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from PIL import ImageDraw
from yolov3_to_onnx import download_file
from data_processing import PreprocessYOLO, PostprocessYOLO, ALL_CATEGORIES
import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import common
TRT_LOGGER = trt.Logger()
def draw_bboxes(image_raw, bboxes, confidences, categories, all_categories, bbox_color='blue'):
    draw = ImageDraw.Draw(image_raw)
    print(bboxes, confidences, categories)
    for box, score, category in zip(bboxes, confidences, categories):
        x_coord, y_coord, width, height = box
        left = max(0, np.floor(x_coord + 0.5).astype(int))
        top = max(0, np.floor(y_coord + 0.5).astype(int))
        right = min(image_raw.width, np.floor(x_coord + width + 0.5).astype(int))
        bottom = min(image_raw.height, np.floor(y_coord + height + 0.5).astype(int))

        draw.rectangle(((left, top), (right, bottom)), outline=bbox_color)
        draw.text((left, top - 12), '{0} {1:.2f}'.format(all_categories[category], score), fill=bbox_color)

    return image_raw
def get_engine(onnx_file_path, engine_file_path=""):
    def build_engine():
        with trt.Builder(TRT_LOGGER) as builder, builder.create_network(common.EXPLICIT_BATCH) as network,  trt.OnnxParser(network, TRT_LOGGER) as parser:
            builder.max_workspace_size = 1<<28
            builder.max_batch_size = 1
            if not os.path.exists(onnx_file_path):
                print(
                    'ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'.format(onnx_file_path))
                exit(0)
            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                if not parser.parse(model.read()):
                    print('ERROR: Failed to parse the ONNX file.')
                    for error in range(parser.num_errors):
                        print(parser.get_error(error))
                    return None
            network.get_input(0).shape = [1,3, 608, 608]
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
            engine = builder.build_cuda_engine(network)
            print("Completed creating Engine")
            with open(engine_file_path, "wb") as f:
                f.write(engine.serialize())
            return engine
    if os.path.exists(engine_file_path):
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()
def main():
    onnx_file_path = 'yolov3.onnx'
    engine_file_path = "yolov3.trt"
    input_image_path = download_file('dog.jpg',
                                     'https://github.com/pjreddie/darknet/raw/f86901f6177dfc6116360a13cc06ab680e0c86b0/data/dog.jpg',
                                     checksum_reference=None)
    input_resolution_yolov3_HW = (608, 608)
    preprocessor = PreprocessYOLO(input_resolution_yolov3_HW)
    image_raw, image = preprocessor.process(input_image_path)
    shape_orig_WH = image_raw.size
    output_shapes = [(1, 255, 19, 19), (1, 255, 38, 38), (1, 255, 76, 76)]
    # Do inference with TensorRT
    trt_outputs = []
    with get_engine(onnx_file_path, engine_file_path) as engine, engine.create_execution_context() as context:
        inputs, outputs, bindings, stream = common.allocate_buffers(engine)
        # Do inference
        print('Running inference on image {}...'.format(input_image_path))
        inputs[0].host = image
        trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
    trt_outputs = [output.reshape(shape) for output, shape in zip(trt_outputs, output_shapes)]

    postprocessor_args = {"yolo_masks": [(6, 7, 8), (3, 4, 5), (0, 1, 2)],
                          # A list of 3 three-dimensional tuples for the YOLO masks
                          "yolo_anchors": [(10, 13), (16, 30), (33, 23), (30, 61), (62, 45),
                                           # A list of 9 two-dimensional tuples for the YOLO anchors
                                           (59, 119), (116, 90), (156, 198), (373, 326)],
                          "obj_threshold": 0.6,  # Threshold for object coverage, float value between 0 and 1
                          "nms_threshold": 0.5,
                          # Threshold for non-max suppression algorithm, float value between 0 and 1
                          "yolo_input_resolution": input_resolution_yolov3_HW}

    postprocessor = PostprocessYOLO(**postprocessor_args)
    boxes, classes, scores = postprocessor.process(trt_outputs, (shape_orig_WH))
    obj_detected_img = draw_bboxes(image_raw, boxes, scores, classes, ALL_CATEGORIES)
    output_image_path = 'dog_bboxes.png'
    obj_detected_img.save(output_image_path, 'PNG')
    print('Saved image with bounding boxes of detected objects to {}.'.format(output_image_path))

if __name__ == '__main__':
    main()