|
|
@@ -1,4 +1,5 @@ |
|
|
|
import torch |
|
|
|
import argparse |
|
|
|
import sys,os |
|
|
|
sys.path.extend(['segutils']) |
|
|
|
from core.models.bisenet import BiSeNet |
|
|
@@ -521,10 +522,12 @@ def EngineInfer(par): |
|
|
|
|
|
|
|
|
|
|
|
if __name__=='__main__': |
|
|
|
pthFile = Path('../weights/road/stdc_360X640.pth') |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
|
|
parser.add_argument('--weights', type=str, default='stdc_360X640.pth', help='model path(s)') |
|
|
|
opt = parser.parse_args() |
|
|
|
print( opt.weights ) |
|
|
|
#pthFile = Path('../../../yolov5TRT/weights/river/stdc_360X640.pth') |
|
|
|
pthFile = Path(opt.weights) |
|
|
|
onnxFile = pthFile.with_suffix('.onnx') |
|
|
|
trtFile = onnxFile.with_suffix('.engine') |
|
|
|
|