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.

711 lines
39KB

  1. # -*- coding: utf-8 -*-
  2. import sys
  3. from enums.ModelTypeEnum import ModelType
  4. from util.GPUtils import get_all_gpu_ids
  5. sys.path.extend(['..', '../AIlib2'])
  6. from AI import AI_process,AI_process_forest,get_postProcess_para
  7. import cv2,os,time
  8. from segutils.segmodel import SegModel
  9. from models.experimental import attempt_load
  10. from utils.torch_utils import select_device
  11. from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images
  12. import numpy as np
  13. import torch,glob
  14. import tensorrt as trt
  15. from utilsK.masterUtils import get_needed_objectsIndex
  16. from utilsK.jkmUtils import pre_process,post_process,get_return_data
  17. # 统一模型入口
  18. class Model:
  19. def __init__(self, device, allowedList=None, logger=None, requestId=None, modeType=None):
  20. try:
  21. logger.info("########################加载{}########################, requestId:{}", modeType.value[2], requestId)
  22. self.modeType = modeType
  23. self.allowedList = allowedList
  24. self.trtFlag_det = True, # 检测模型是否采用TRT
  25. self.trtFlag_seg = True, # 分割模型是否采用TRT
  26. if self.trtFlag_det and self.trtFlag_seg:
  27. gpu = get_all_gpu_ids()[int(device)]
  28. if '3090' in gpu.name:
  29. detweights = "../AIlib2/weights/%s/yolov5_3090_fp16.engine" % modeType.value[3]
  30. segweights = '../AIlib2/weights/%s/stdc_360X640_3090_fp16.engine' % modeType.value[3]
  31. elif '2080' in gpu.name:
  32. detweights = "../AIlib2/weights/%s/yolov5_2080Ti_fp16.engine" % modeType.value[3]
  33. segweights = '../AIlib2/weights/%s/stdc_360X640_2080Ti_fp16.engine' % modeType.value[3]
  34. elif '4090' in gpu.name:
  35. detweights = "../AIlib2/weights/%s/yolov5_4090_fp16.engine" % modeType.value[3]
  36. segweights = '../AIlib2/weights/%s/stdc_360X640_4090_fp16.engine' % modeType.value[3]
  37. elif 'A10' in gpu.name:
  38. detweights = "../AIlib2/weights/%s/yolov5_A10_fp16.engine" % modeType.value[3]
  39. segweights = '../AIlib2/weights/%s/stdc_360X640_A10_fp16.engine' % modeType.value[3]
  40. else:
  41. raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  42. else:
  43. detweights = "../AIlib2/weights/conf/%s/yolov5.pt" % modeType.value[3]
  44. segweights = '../AIlib2/weights/conf/%s/stdc_360X640.pth' % modeType.value[3]
  45. par = {
  46. 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  47. 'labelnames': "../AIlib2/weights/conf/%s/labelnames.json" % modeType.value[3], # 检测类别对照表
  48. 'trtFlag_det': self.trtFlag_det, # 检测模型是否采用TRT
  49. 'trtFlag_seg': self.trtFlag_seg, # 分割模型是否采用TRT
  50. 'Detweights': detweights, # 检测模型路径
  51. 'seg_nclass': 2, # 分割模型类别数目,默认2类
  52. 'segPar': {
  53. 'modelSize': (640, 360),
  54. 'mean': (0.485, 0.456, 0.406),
  55. 'std' : (0.229, 0.224, 0.225),
  56. 'numpy': False,
  57. 'RGB_convert_first': True
  58. }, # 分割模型预处理参数
  59. 'Segweights': segweights, # 分割模型权重位置
  60. 'postFile': '../AIlib2/weights/conf/%s/para.json' % modeType.value[3] # 后处理参数文件
  61. }
  62. if ModelType.WATER_SURFACE_MODEL == modeType:
  63. par['segRegionCnt'] = 1
  64. par['slopeIndex'] = [5, 6, 7]
  65. if ModelType.TRAFFIC_FARM_MODEL == modeType:
  66. par['segRegionCnt'] = 2
  67. par['slopeIndex'] = [5, 6, 7]
  68. else:
  69. par['segRegionCnt'] = 0
  70. par['slopeIndex'] = []
  71. par['segPar'] = None
  72. par['Segweights'] = None
  73. self.device = select_device(par.get('device'))
  74. self.names = get_labelnames(par.get('labelnames'))
  75. self.half = self.device.type != 'cpu'
  76. Detweights = par.get('Detweights') # 升级后的检测模型
  77. if self.trtFlag_det:
  78. with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
  79. self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
  80. print('####load TRT model :%s'%(Detweights))
  81. else:
  82. self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  83. if self.half:
  84. self.model.half()
  85. seg_nclass = par.get('seg_nclass')
  86. self.segPar = par.get('segPar')
  87. Segweights = par.get('Segweights')
  88. if Segweights is not None:
  89. if self.trtFlag_seg:
  90. log = trt.Logger(trt.Logger.ERROR)
  91. with open(Segweights, "rb") as f, trt.Runtime(log) as runtime:
  92. self.segmodel = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件,返回ICudaEngine对象
  93. print('############locad seg model trt success#######')
  94. else:
  95. self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device)
  96. print('############locad seg model pth success#######')
  97. else:
  98. self.segmodel = None
  99. self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  100. self.objectPar = {
  101. 'half': self.half,
  102. 'device': self.device,
  103. 'conf_thres': self.conf_thres,
  104. 'iou_thres': self.iou_thres,
  105. 'allowedList': self.allowedList,
  106. 'slopeIndex': par.get('slopeIndex'),
  107. 'segRegionCnt': par.get('segRegionCnt'),
  108. 'trtFlag_det': self.trtFlag_det,
  109. 'trtFlag_seg': self.trtFlag_seg
  110. }
  111. self.label_arraylist = None
  112. self.digitFont = None
  113. except Exception as ee:
  114. logger.exception("异常:{}, requestId:{}", ee, requestId)
  115. raise ee
  116. def process(self, frame, width=1920):
  117. if self.label_arraylist is None:
  118. fontsize = int(width / 1920 * 40)
  119. line_thickness = 1
  120. boxLine_thickness = 1
  121. waterLineWidth = 1
  122. if width >= 960:
  123. line_thickness = int(round(width / 1920 * 3) - 1)
  124. boxLine_thickness = int(round(width / 1920 * 3))
  125. waterLineWidth = int(round(width / 1920 * 3))
  126. numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小
  127. self.digitFont = {'line_thickness': line_thickness,
  128. 'boxLine_thickness': boxLine_thickness,
  129. 'fontSize': numFontSize,
  130. 'waterLineColor': (0, 255, 255),
  131. 'waterLineWidth': waterLineWidth}
  132. self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  133. fontpath="../AIlib2/conf/platech.ttf")
  134. if ModelType.WATER_SURFACE_MODEL == self.modeType or ModelType.TRAFFIC_FARM_MODEL == self.modeType:
  135. return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  136. self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar)
  137. else:
  138. return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  139. self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,
  140. self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)
  141. # # 河道模型
  142. # class SZModel:
  143. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  144. # try:
  145. # logger.info("########################加载河道模型########################, requestId:{}", requestId)
  146. # self.allowedList = allowedList
  147. # trtFlag_det = False, # 检测模型是否采用TRT
  148. # trtFlag_seg = False, # 分割模型是否采用TRT
  149. # if trtFlag_det and trtFlag_seg:
  150. # gpu = get_all_gpu_ids()[int(device)]
  151. # if '3090' in gpu.name:
  152. # det_gpuname = '_3090_fp16.engine'
  153. # seg_gpuname = '_3090_fp16.engine'
  154. # elif '2080' in gpu.name:
  155. # det_gpuname = '_2080Ti_fp16.engine'
  156. # seg_gpuname = '_2080Ti_fp16.engine'
  157. # elif '4090' in gpu.name:
  158. # det_gpuname = '_4090_fp16.engine'
  159. # seg_gpuname = '_4090_fp16.engine'
  160. # elif 'A10' in gpu.name:
  161. # det_gpuname = '_A10_fp16.engine'
  162. # seg_gpuname = '_A10_fp16.engine'
  163. # else:
  164. # raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  165. # else:
  166. # det_gpuname = '.pt'
  167. # seg_gpuname = '.pth'
  168. # par = {
  169. # 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  170. # 'labelnames': "../AIlib2/weights/river/labelnames.json", # 检测类别对照表
  171. # 'trtFlag_det': trtFlag_det, # 检测模型是否采用TRT
  172. # 'trtFlag_seg': trtFlag_seg, # 分割模型是否采用TRT
  173. # 'Detweights': "../AIlib2/weights/river/yolov5%s" % det_gpuname, # 检测模型路径
  174. # 'slopeIndex': [5, 6, 7], # 岸坡类别(或者其它业务里的类别),不与河道(分割的前景区域)计算交并比,即不论是否在河道内都显示。
  175. # 'seg_nclass': 2, # 分割模型类别数目,默认2类
  176. # 'segRegionCnt': 1, # 分割模型结果需要保留的等值线数目
  177. # 'segPar': {
  178. # 'modelSize': (640,360),
  179. # 'mean': (0.485, 0.456, 0.406),
  180. # 'std' : (0.229, 0.224, 0.225),
  181. # 'numpy': False,
  182. # 'RGB_convert_first': True
  183. # }, # 分割模型预处理参数
  184. # 'Segweights': '../AIlib2/weights/river/stdc_360X640%s' % seg_gpuname, # 分割模型权重位置
  185. # 'postFile': '../AIlib2/weights/river/para.json' # 后处理参数文件
  186. # }
  187. # self.device = select_device(par.get('device'))
  188. # self.names = get_labelnames(par.get('labelnames'))
  189. # half = self.device.type != 'cpu'
  190. # Detweights = par.get('Detweights') # 升级后的检测模型
  191. # if trtFlag_det:
  192. # log = trt.Logger(trt.Logger.ERROR)
  193. # with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:
  194. # self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
  195. # print('############locad det model trt success#######')
  196. # else:
  197. # self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  198. # print('############locad det model pth success#######')
  199. # if half:
  200. # self.model.half()
  201. # seg_nclass = par.get('seg_nclass')
  202. # self.segPar = par.get('segPar')
  203. # Segweights = par.get('Segweights')
  204. # if trtFlag_seg:
  205. # log = trt.Logger(trt.Logger.ERROR)
  206. # with open(Segweights, "rb") as f, trt.Runtime(log) as runtime:
  207. # self.segmodel = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件,返回ICudaEngine对象
  208. # print('############locad seg model trt success#######')
  209. # else:
  210. # self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device)
  211. # print('############locad seg model pth success#######')
  212. #
  213. # self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  214. # self.objectPar = {
  215. # 'half': half,
  216. # 'device': self.device,
  217. # 'conf_thres': self.conf_thres,
  218. # 'iou_thres': self.iou_thres,
  219. # 'allowedList': self.allowedList,
  220. # 'slopeIndex': par.get('slopeIndex'),
  221. # 'segRegionCnt': par.get('segRegionCnt'),
  222. # 'trtFlag_det': trtFlag_det,
  223. # 'trtFlag_seg': trtFlag_seg
  224. # }
  225. # self.label_arraylist = None
  226. # self.digitFont = None
  227. # except Exception as ee:
  228. # logger.exception("异常:{}, requestId:{}", ee, requestId)
  229. #
  230. #
  231. #
  232. # # names, label_arraylist, rainbows, conf_thres, iou_thres
  233. # def process(self, frame, width=1920):
  234. # if self.label_arraylist is None:
  235. # fontsize = int(width / 1920 * 40)
  236. # line_thickness = 1
  237. # boxLine_thickness = 1
  238. # waterLineWidth = 1
  239. # if width >= 960:
  240. # line_thickness = int(round(width / 1920 * 3) - 1)
  241. # boxLine_thickness = int(round(width / 1920 * 3))
  242. # waterLineWidth = int(round(width / 1920 * 3))
  243. # numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小
  244. # self.digitFont = {'line_thickness': line_thickness,
  245. # 'boxLine_thickness': boxLine_thickness,
  246. # 'fontSize': numFontSize,
  247. # 'waterLineColor': (0, 255, 255),
  248. # 'waterLineWidth': waterLineWidth}
  249. # self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  250. # fontpath="../AIlib2/conf/platech.ttf")
  251. # return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  252. # self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar)
  253. #
  254. #
  255. # # 森林模型
  256. # class LCModel:
  257. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  258. # logger.info("########################加载森林模型########################, requestId:{}", requestId)
  259. # self.allowedList = allowedList
  260. # self.trtFlag_det = True, # 检测模型是否采用TRT
  261. # if self.trtFlag_det:
  262. # gpu = get_all_gpu_ids()[int(device)]
  263. # if '3090' in gpu.name:
  264. # det_gpuname = '_3090_fp16.engine'
  265. # elif '2080' in gpu.name:
  266. # det_gpuname = '_2080Ti_fp16.engine'
  267. # elif '4090' in gpu.name:
  268. # det_gpuname = '_4090_fp16.engine'
  269. # elif 'A10' in gpu.name:
  270. # det_gpuname = '_A10_fp16.engine'
  271. # else:
  272. # raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  273. # else:
  274. # det_gpuname = '.pt'
  275. # par = {
  276. # 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  277. # 'labelnames': "../AIlib2/weights/forest/labelnames.json", # 检测类别对照表
  278. # 'Detweights': "../AIlib2/weights/forest/yolov5%s" % det_gpuname, # 检测模型路径
  279. # 'slopeIndex': [], # 岸坡类别(或者其它业务里的类别),不与河道(分割的前景区域)计算交并比,即不论是否在河道内都显示。
  280. # 'seg_nclass': 2, # 分割模型类别数目,默认2类
  281. # 'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目
  282. # 'segPar': None, # 分割模型预处理参数
  283. # 'Segweights': None, # 分割模型权重位置
  284. # 'postFile': '../AIlib2/weights/forest/para.json' # 后处理参数文件
  285. # }
  286. # self.device = select_device(par.get('device'))
  287. # self.half = self.device.type != 'cpu' # half precision only supported on CUDA
  288. # Detweights = par.get('Detweights')
  289. # if self.trtFlag_det:
  290. # log = trt.Logger(trt.Logger.ERROR)
  291. # with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:
  292. # self.model=runtime.deserialize_cuda_engine(f.read())
  293. # print('####load TRT model :%s'%(Detweights))
  294. # else:
  295. # self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  296. # if self.half:
  297. # self.model.half()
  298. # self.names = get_labelnames(par.get('labelnames'))
  299. # self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  300. # self.digitFont = None
  301. # self.label_arraylist = None
  302. # self.segmodel = None
  303. #
  304. # # names, label_arraylist, rainbows, conf_thres, iou_thres
  305. # def process(self, frame, width=1920):
  306. # if self.label_arraylist is None:
  307. # fontsize = int(width / 1920 * 40)
  308. # line_thickness = 1
  309. # boxLine_thickness = 1
  310. # waterLineWidth = 1
  311. # if width >= 960:
  312. # line_thickness = int(round(width / 1920 * 3) - 1)
  313. # boxLine_thickness = int(round(width / 1920 * 3))
  314. # waterLineWidth = int(round(width / 1920 * 3))
  315. # numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小
  316. # self.digitFont = {'line_thickness': line_thickness,
  317. # 'boxLine_thickness': boxLine_thickness,
  318. # 'fontSize': numFontSize,
  319. # 'waterLineColor': (0, 255, 255),
  320. # 'waterLineWidth': waterLineWidth}
  321. # self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  322. # fontpath="../AIlib2/conf/platech.ttf")
  323. # return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  324. # self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,
  325. # self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)
  326. #
  327. #
  328. # # 车辆模型
  329. # class VehicleModel:
  330. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  331. # logger.info("########################加载车辆模型########################, requestId:{}", requestId)
  332. # self.allowedList = allowedList
  333. # self.trtFlag_det = True, # 检测模型是否采用TRT
  334. # if self.trtFlag_det:
  335. # gpu = get_all_gpu_ids()[int(device)]
  336. # if '3090' in gpu.name:
  337. # det_gpuname = '_3090_fp16.engine'
  338. # elif '2080' in gpu.name:
  339. # det_gpuname = '_2080Ti_fp16.engine'
  340. # elif '4090' in gpu.name:
  341. # det_gpuname = '_4090_fp16.engine'
  342. # elif 'A10' in gpu.name:
  343. # det_gpuname = '_A10_fp16.engine'
  344. # else:
  345. # raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  346. # else:
  347. # det_gpuname = '.pt'
  348. # par = {
  349. # 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  350. # 'labelnames': "../AIlib2/weights/vehicle/labelnames.json", # 检测类别对照表
  351. # 'Detweights': "../AIlib2/weights/vehicle/yolov5%s" % det_gpuname, # 检测模型路径
  352. # 'slopeIndex': [], # 岸坡类别(或者其它业务里的类别),不与河道(分割的前景区域)计算交并比,即不论是否在河道内都显示。
  353. # 'seg_nclass': 2, # 分割模型类别数目,默认2类
  354. # 'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目
  355. # 'segPar': None, # 分割模型预处理参数
  356. # 'Segweights': None, # 分割模型权重位置
  357. # 'postFile': '../AIlib2/weights/vehicle/para.json' # 后处理参数文件
  358. # }
  359. # self.device = select_device(par.get('device'))
  360. # self.half = self.device.type != 'cpu' # half precision only supported on CUDA
  361. # Detweights = par.get('Detweights')
  362. # if self.trtFlag_det:
  363. # log = trt.Logger(trt.Logger.ERROR)
  364. # with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:
  365. # self.model=runtime.deserialize_cuda_engine(f.read())
  366. # print('####load TRT model :%s'%(Detweights))
  367. # else:
  368. # self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  369. # if self.half:
  370. # self.model.half()
  371. # self.names = get_labelnames(par.get('labelnames'))
  372. # self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  373. # self.digitFont = None
  374. # self.label_arraylist = None
  375. # self.segmodel = None
  376. #
  377. # # names, label_arraylist, rainbows, conf_thres, iou_thres
  378. # def process(self, frame, width=1920):
  379. # if self.label_arraylist is None:
  380. # fontsize = int(width / 1920 * 40)
  381. # line_thickness = 1
  382. # boxLine_thickness = 1
  383. # waterLineWidth = 1
  384. # if width >= 960:
  385. # line_thickness = int(round(width / 1920 * 3) - 1)
  386. # boxLine_thickness = int(round(width / 1920 * 3))
  387. # waterLineWidth = int(round(width / 1920 * 3))
  388. # numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小
  389. # self.digitFont = {'line_thickness': line_thickness,
  390. # 'boxLine_thickness': boxLine_thickness,
  391. # 'fontSize': numFontSize,
  392. # 'waterLineColor': (0, 255, 255),
  393. # 'waterLineWidth': waterLineWidth}
  394. # self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  395. # fontpath="../AIlib2/conf/platech.ttf")
  396. # return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  397. # self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,
  398. # self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)
  399. # # 行人模型
  400. # class PedestrianModel:
  401. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  402. # logger.info("########################加载行人模型########################, requestId:{}", requestId)
  403. # self.allowedList = allowedList
  404. # self.trtFlag_det = True, # 检测模型是否采用TRT
  405. # if self.trtFlag_det:
  406. # gpu = get_all_gpu_ids()[int(device)]
  407. # if '3090' in gpu.name:
  408. # det_gpuname = '_3090_fp16.engine'
  409. # elif '2080' in gpu.name:
  410. # det_gpuname = '_2080Ti_fp16.engine'
  411. # elif '4090' in gpu.name:
  412. # det_gpuname = '_4090_fp16.engine'
  413. # elif 'A10' in gpu.name:
  414. # det_gpuname = '_A10_fp16.engine'
  415. # else:
  416. # raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  417. # else:
  418. # det_gpuname = '.pt'
  419. # par = {
  420. # 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  421. # 'labelnames': "../AIlib2/weights/pedestrian/labelnames.json", # 检测类别对照表
  422. # 'Detweights': "../AIlib2/weights/pedestrian/yolov5%s" % det_gpuname, # 检测模型路径
  423. # 'slopeIndex': [], # 岸坡类别(或者其它业务里的类别),不与河道(分割的前景区域)计算交并比,即不论是否在河道内都显示。
  424. # 'seg_nclass': 2, # 分割模型类别数目,默认2类
  425. # 'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目
  426. # 'segPar': None, # 分割模型预处理参数
  427. # 'Segweights': None, # 分割模型权重位置
  428. # 'postFile': '../AIlib2/weights/pedestrian/para.json' # 后处理参数文件
  429. # }
  430. # self.device = select_device(par.get('device'))
  431. # self.half = self.device.type != 'cpu' # half precision only supported on CUDA
  432. # Detweights = par.get('Detweights')
  433. # if self.trtFlag_det:
  434. # log = trt.Logger(trt.Logger.ERROR)
  435. # with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:
  436. # self.model=runtime.deserialize_cuda_engine(f.read())
  437. # print('####load TRT model :%s'%(Detweights))
  438. # else:
  439. # self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  440. # if self.half:
  441. # self.model.half()
  442. # self.names = get_labelnames(par.get('labelnames'))
  443. # self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  444. # self.digitFont = None
  445. # self.label_arraylist = None
  446. # self.segmodel = None
  447. #
  448. # # names, label_arraylist, rainbows, conf_thres, iou_thres
  449. # def process(self, frame, width=1920):
  450. # if self.label_arraylist is None:
  451. # fontsize = int(width / 1920 * 40)
  452. # line_thickness = 1
  453. # boxLine_thickness = 1
  454. # waterLineWidth = 1
  455. # if width >= 960:
  456. # line_thickness = int(round(width / 1920 * 3) - 1)
  457. # boxLine_thickness = int(round(width / 1920 * 3))
  458. # waterLineWidth = int(round(width / 1920 * 3))
  459. # numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小
  460. # self.digitFont = {'line_thickness': line_thickness,
  461. # 'boxLine_thickness': boxLine_thickness,
  462. # 'fontSize': numFontSize,
  463. # 'waterLineColor': (0, 255, 255),
  464. # 'waterLineWidth': waterLineWidth}
  465. # self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  466. # fontpath="../AIlib2/conf/platech.ttf")
  467. # return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  468. # self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,
  469. # self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)
  470. #
  471. #
  472. # # 烟火模型
  473. # class SmogfireModel:
  474. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  475. # logger.info("########################加载烟火模型########################, requestId:{}", requestId)
  476. # self.allowedList = allowedList
  477. # self.trtFlag_det = True, # 检测模型是否采用TRT
  478. # if self.trtFlag_det:
  479. # gpu = get_all_gpu_ids()[int(device)]
  480. # if '3090' in gpu.name:
  481. # det_gpuname = '_3090_fp16.engine'
  482. # elif '2080' in gpu.name:
  483. # det_gpuname = '_2080Ti_fp16.engine'
  484. # elif '4090' in gpu.name:
  485. # det_gpuname = '_4090_fp16.engine'
  486. # elif 'A10' in gpu.name:
  487. # det_gpuname = '_A10_fp16.engine'
  488. # else:
  489. # raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  490. # else:
  491. # det_gpuname = '.pt'
  492. # par = {
  493. # 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  494. # 'labelnames': "../AIlib2/weights/smogfire/labelnames.json", # 检测类别对照表
  495. # 'Detweights': "../AIlib2/weights/smogfire/yolov5%s" % det_gpuname, # 检测模型路径
  496. # 'slopeIndex': [], # 岸坡类别(或者其它业务里的类别),不与河道(分割的前景区域)计算交并比,即不论是否在河道内都显示。
  497. # 'seg_nclass': 2, # 分割模型类别数目,默认2类
  498. # 'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目
  499. # 'segPar': None, # 分割模型预处理参数
  500. # 'Segweights': None, # 分割模型权重位置
  501. # 'postFile': '../AIlib2/weights/smogfire/para.json' # 后处理参数文件
  502. # }
  503. # self.device = select_device(par.get('device'))
  504. # self.half = self.device.type != 'cpu' # half precision only supported on CUDA
  505. # Detweights = par.get('Detweights')
  506. # if self.trtFlag_det:
  507. # log = trt.Logger(trt.Logger.ERROR)
  508. # with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:
  509. # self.model=runtime.deserialize_cuda_engine(f.read())
  510. # print('####load TRT model :%s'%(Detweights))
  511. # else:
  512. # self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  513. # if self.half:
  514. # self.model.half()
  515. # self.names = get_labelnames(par.get('labelnames'))
  516. # self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  517. # self.digitFont = None
  518. # self.label_arraylist = None
  519. # self.segmodel = None
  520. #
  521. # # names, label_arraylist, rainbows, conf_thres, iou_thres
  522. # def process(self, frame, width=1920):
  523. # if self.label_arraylist is None:
  524. # fontsize = int(width / 1920 * 40)
  525. # line_thickness = 1
  526. # boxLine_thickness = 1
  527. # waterLineWidth = 1
  528. # if width >= 960:
  529. # line_thickness = int(round(width / 1920 * 3) - 1)
  530. # boxLine_thickness = int(round(width / 1920 * 3))
  531. # waterLineWidth = int(round(width / 1920 * 3))
  532. # numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小
  533. # self.digitFont = {'line_thickness': line_thickness,
  534. # 'boxLine_thickness': boxLine_thickness,
  535. # 'fontSize': numFontSize,
  536. # 'waterLineColor': (0, 255, 255),
  537. # 'waterLineWidth': waterLineWidth}
  538. # self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  539. # fontpath="../AIlib2/conf/platech.ttf")
  540. # return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  541. # self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,
  542. # self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)
  543. #
  544. # # 交通模型
  545. # class RFModel:
  546. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  547. # logger.info("########################加载交通模型########################, requestId:{}", requestId)
  548. # self.allowedList = allowedList
  549. # trtFlag_det = True, # 检测模型是否采用TRT
  550. # trtFlag_seg = True, # 分割模型是否采用TRT
  551. # if trtFlag_det and trtFlag_seg:
  552. # gpu = get_all_gpu_ids()[int(device)]
  553. # if '3090' in gpu.name:
  554. # det_gpuname = '_3090_fp16.engine'
  555. # seg_gpuname = '_3090_fp16.engine'
  556. # elif '2080' in gpu.name:
  557. # det_gpuname = '_2080Ti_fp16.engine'
  558. # seg_gpuname = '_2080Ti_fp16.engine'
  559. # elif '4090' in gpu.name:
  560. # det_gpuname = '_4090_fp16.engine'
  561. # seg_gpuname = '_4090_fp16.engine'
  562. # elif 'A10' in gpu.name:
  563. # det_gpuname = '_A10_fp16.engine'
  564. # seg_gpuname = '_A10_fp16.engine'
  565. # else:
  566. # raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)
  567. # else:
  568. # det_gpuname = '.pt'
  569. # seg_gpuname = '.pth'
  570. # par = {
  571. # 'device': device, # 显卡号,如果用TRT模型,只支持0(单显卡)
  572. # 'labelnames': "../AIlib2/weights/road/labelnames.json", # 检测类别对照表
  573. # 'trtFlag_det': trtFlag_det, # 检测模型是否采用TRT
  574. # 'trtFlag_seg': trtFlag_seg, # 分割模型是否采用TRT
  575. # 'Detweights': "../AIlib2/weights/road/yolov5%s" % det_gpuname, # 检测模型路径
  576. # 'slopeIndex': [], # 岸坡类别(或者其它业务里的类别),不与河道(分割的前景区域)计算交并比,即不论是否在河道内都显示。
  577. # 'seg_nclass': 2, # 分割模型类别数目,默认2类
  578. # 'segRegionCnt': 2, # 分割模型结果需要保留的等值线数目
  579. # 'segPar': {
  580. # 'modelSize': (640,360),
  581. # 'mean': (0.485, 0.456, 0.406),
  582. # 'std' : (0.229, 0.224, 0.225),
  583. # 'numpy': False,
  584. # 'RGB_convert_first': True
  585. # }, # 分割模型预处理参数
  586. # 'Segweights': '../AIlib2/weights/road/stdc_360X640%s' % seg_gpuname, # 分割模型权重位置
  587. # 'postFile': '../AIlib2/weights/road/para.json' # 后处理参数文件
  588. # }
  589. # self.device = select_device(par.get('device'))
  590. # self.names = get_labelnames(par.get('labelnames'))
  591. # half = self.device.type != 'cpu'
  592. # Detweights = par.get('Detweights') # 升级后的检测模型
  593. # if trtFlag_det:
  594. # log = trt.Logger(trt.Logger.ERROR)
  595. # with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:
  596. # self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
  597. # print('############locad det model trt success#######')
  598. # else:
  599. # self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model
  600. # print('############locad det model pth success#######')
  601. # if half:
  602. # self.model.half()
  603. # seg_nclass = par.get('seg_nclass')
  604. # self.segPar = par.get('segPar')
  605. # Segweights = par.get('Segweights')
  606. # if trtFlag_seg:
  607. # log = trt.Logger(trt.Logger.ERROR)
  608. # with open(Segweights, "rb") as f, trt.Runtime(log) as runtime:
  609. # self.segmodel = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件,返回ICudaEngine对象
  610. # print('############locad seg model trt success#######')
  611. # else:
  612. # self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device)
  613. # print('############locad seg model pth success#######')
  614. #
  615. # self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))
  616. # self.objectPar = {
  617. # 'half': half,
  618. # 'device': self.device,
  619. # 'conf_thres': self.conf_thres,
  620. # 'iou_thres': self.iou_thres,
  621. # 'allowedList': self.allowedList,
  622. # 'slopeIndex': par.get('slopeIndex'),
  623. # 'segRegionCnt': par.get('segRegionCnt'),
  624. # 'trtFlag_det': trtFlag_det,
  625. # 'trtFlag_seg': trtFlag_seg
  626. # }
  627. # self.label_arraylist = None
  628. # self.digitFont = None
  629. #
  630. # def process(self, frame, width=1920):
  631. # if self.label_arraylist is None:
  632. # fontsize = int(width / 1920 * 40)
  633. # line_thickness = 1
  634. # boxLine_thickness = 1
  635. # waterLineWidth = 1
  636. # if width >= 960:
  637. # line_thickness = int(round(width / 1920 * 3) - 1)
  638. # boxLine_thickness = int(round(width / 1920 * 2))
  639. # waterLineWidth = int(round(width / 1920 * 3))
  640. # numFontSize = float(format(width / 1920 * 1.1, '.1f'))
  641. #
  642. # self.digitFont = {'line_thickness': line_thickness,
  643. # 'boxLine_thickness': boxLine_thickness,
  644. # 'fontSize': numFontSize,
  645. # 'waterLineColor': (0, 255, 255),
  646. # 'waterLineWidth': waterLineWidth}
  647. # self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,
  648. # fontpath="../AIlib2/conf/platech.ttf")
  649. # return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist,
  650. # self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar)
  651. # class IMModel:
  652. # def __init__(self, device, allowedList=None, logger=None, requestId=None):
  653. # logger.info("########################加载防疫模型########################, requestId:{}", requestId)
  654. # self.allowedList = allowedList
  655. # self.img_type = 'code'
  656. # self.par = {'code': {'weights': '../AIlib2/weights/jkm/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10},
  657. # 'plate': {'weights': '../AIlib2/weights/jkm/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1},
  658. # 'conf_thres': 0.4,
  659. # 'iou_thres': 0.45,
  660. # 'device': 'cuda:%s' % device,
  661. # 'plate_dilate': (0.5, 0.3)
  662. # }
  663. # self.device = torch.device(self.par['device'])
  664. # self.model = torch.jit.load(self.par[self.img_type]['weights'])
  665. #
  666. # def process(self, frame):
  667. # img, padInfos = pre_process(frame, self.device) ##预处理
  668. # pred = self.model(img) ##模型推理
  669. # boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'],
  670. # iou_thres=self.par['iou_thres'], nc=self.par[self.img_type]['nc']) # 后处理
  671. # dataBack = get_return_data(frame, boxes, modelType=self.img_type, plate_dilate=self.par['plate_dilate'])
  672. # return dataBack
  673. # 车牌分割模型、健康码、行程码分割模型
  674. class IMModel:
  675. def __init__(self, device, allowedList=None, logger=None, requestId=None, modeType=None):
  676. logger.info("########################加载{}########################, requestId:{}", modeType.value[2], requestId)
  677. self.allowedList = allowedList
  678. if ModelType.EPIDEMIC_PREVENTION_MODEL == modeType:
  679. self.img_type = 'code' ## code,plate
  680. if ModelType.PLATE_MODEL == modeType:
  681. self.img_type = 'plate' ## code,plate
  682. self.par = {'code': {'weights': '../AIlib2/weights/conf/jkm/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10},
  683. 'plate': {'weights': '../AIlib2/weights/conf/jkm/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1},
  684. 'conf_thres': 0.4,
  685. 'iou_thres': 0.45,
  686. 'device': 'cuda:%s' % device,
  687. 'plate_dilate': (0.5, 0.3)
  688. }
  689. self.device = torch.device(self.par['device'])
  690. self.model = torch.jit.load(self.par[self.img_type]['weights'])
  691. def process(self, frame):
  692. img, padInfos = pre_process(frame, self.device) ##预处理
  693. pred = self.model(img) ##模型推理
  694. boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'],
  695. iou_thres=self.par['iou_thres'], nc=self.par[self.img_type]['nc']) # 后处理
  696. dataBack = get_return_data(frame, boxes, modelType=self.img_type, plate_dilate=self.par['plate_dilate'])
  697. return dataBack