You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

48 lines
1.6KB

  1. import os
  2. import time,argparse
  3. import cv2
  4. import torch
  5. import sys
  6. sys.path.extend(['..' ])
  7. from DMPRUtils.model.detector import DirectionalPointDetector
  8. from pathlib import Path
  9. from segutils.trtUtils import toONNX,ONNXtoTrt
  10. from DMPRUtils.yolo_net import Model
  11. def main(opt):
  12. pars={'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.3, 'dmprimg_size':640,
  13. 'mWidth':640,'mHeight':640
  14. }
  15. ##以下参数目前不可改
  16. #DMPRweights = "weights/urbanManagement/DMPR/dp_detector_499.pth"
  17. DMPRweights = opt.weights.strip()
  18. DMPR_pthFile = Path(DMPRweights)
  19. inputShape =(1, 3, pars['mHeight'],pars['mWidth'])#(bs,channels,height,width)
  20. DMPR_onnxFile = str(DMPR_pthFile.with_suffix('.onnx'))
  21. DMPR_trtFile = DMPR_onnxFile.replace('.onnx','.engine' )
  22. ##加载模型,准备好显示字符
  23. device = 'cuda:0'
  24. # DMPR model
  25. #DMPRmodel = DirectionalPointDetector(3, pars['depth_factor'], pars['NUM_FEATURE_MAP_CHANNEL']).to(device)
  26. confUrl = os.path.join( os.path.dirname(__file__),'config','yolov5s.yaml' )
  27. DMPRmodel = Model(confUrl, ch=3).to(device)
  28. DMPRmodel.load_state_dict(torch.load(DMPRweights))
  29. toONNX(DMPRmodel,DMPR_onnxFile,inputShape=inputShape,device=device,dynamic=True)
  30. ONNXtoTrt(DMPR_onnxFile,DMPR_trtFile)
  31. if __name__ == '__main__':
  32. parser = argparse.ArgumentParser()
  33. parser.add_argument('--weights', type=str, default='/mnt/thsw2/DSP2/weights/cityMangement2/weights/urbanManagement/DMPR/dp_detector_499.pth', help='model path(s)')
  34. opt = parser.parse_args()
  35. main(opt)