classifier_trt_image.py

import utils.inference as inference_utils # TRT/TF inference wrappers
import utils.model as model_utils # UFF conversion
import tensorrt as trt
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='TRT params')
    parser.add_argument('--FP', default="32", type=str)
    args = parser.parse_args()
    FP = args.FP

    if FP == '32':
        trt_datatype = trt.DataType.FLOAT
        engine_name = 'resnet50_engine.buf'
    elif FP == '16':
        trt_datatype = trt.DataType.HALF
        engine_name = 'resnet50_engine_half.buf'
    else:
        print('Wrong precision')
    onnx_model_path = 'resnet50_simple.onnx'
    input_img_path =  'hotdog.jpg'

    trt_inference_wrapper = inference_utils.TRTInference(
        engine_name, onnx_model_path, trt_datatype)
    print('done, running inference')
    detection_out = trt_inference_wrapper.infer(input_img_path, 'image')
    print(detection_out)
    print('finished inference')