labels = ['can', 'paper_cup', 'paper_box', 'paper_milkbox', 'plastic'] print("[*] Importing packages...") import common import tensorrt as trt import os import cv2 import pandas as pd import numpy as np TRT_LOGGER = trt.Logger(trt.Logger.WARNING) print("[*] Loading model...") # load trt engine trt_path = 'model.trt' with open(trt_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) inputs, outputs, bindings, stream = common.allocate_buffers(engine) if __name__ == '__main__': def pred(f, dirpath): img = cv2.imread(os.path.join(dirpath, f)) weight = df.loc[f]['weight'] inputs[0].host = np.expand_dims(img, 0).astype('float32') inputs[1].host = np.expand_dims(weight, 0).astype('float32') # inference with engine.create_execution_context() as context: trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs,stream=stream) result = trt_outputs[0].argmax(-1) return labels[result] df = pd.read_csv('test_data/weights_test.csv') df = df.set_index('name') for dirpath, dirnames, filenames in os.walk('test_data'): for f in filenames: if f.endswith('.jpg'): print(f'{f}: {pred(f, dirpath)}')