JardinRyu / Jetson_Nano_People_Counting

MobileNetV1-SSD + SORT based Real-Time Tracking and Counting on Jetson Nano
111 stars 32 forks source link

If I use my own yolov5 detecotr #3

Open fanbbbb opened 3 years ago

fanbbbb commented 3 years ago

I have run this repo with SSD mobilenet successfully. Now I want to use my own yolov5 tensorrt model. I changed utils/ssd.py line 49-57 to

def _load_plugins(self):
        if trt.__version__[0] < '7':
            # ctypes.CDLL("ssd/libflattenconcat.so")
            PLUGIN_LIBRARY = "yolov5/libmyplugins.so"
            ctypes.CDLL(PLUGIN_LIBRARY)
        trt.init_libnvinfer_plugins(self.trt_logger, '')
def _load_engine(self):
        TRTbin = "yolov5/best.engine"
        # TRTbin = 'ssd/TRT_%s.bin' % self.model
        with open(TRTbin, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
            return runtime.deserialize_cuda_engine(f.read())

Should I change something else to make sure the code can run?