gitadmin
/
tuoheng_alg_move_20240715
Observar 1
Favorito 0
Fork 0
Código Issues 0 Pull requests 0 Versões 11 Wiki Atividade
Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
60 Commits
8 Branches
Tag: 6c0f5c94bc
Branches Tags
${ item.name }
Criar branch ${ searchTerm }
de 6c0f5c94bc
${ noResults }
tuoheng_alg_move_20240715/util/ModelUtils.py

711 linhas
39KB
Original Anotar Histórico 

  1. # -*- coding: utf-8 -*-

  2. import sys

  3. 

  4. from enums.ModelTypeEnum import ModelType

  5. from util.GPUtils import get_all_gpu_ids

  6. 

  7. sys.path.extend(['..', '../AIlib2'])

  8. from AI import AI_process,AI_process_forest,get_postProcess_para

  9. import cv2,os,time

  10. from segutils.segmodel import SegModel

  11. from models.experimental import attempt_load

  12. from utils.torch_utils import select_device

  13. from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images

  14. import numpy as np

  15. import torch,glob

  16. import tensorrt as trt

  17. from utilsK.masterUtils import get_needed_objectsIndex

  18. from utilsK.jkmUtils import pre_process,post_process,get_return_data

  19. 

  20. 

  21. # 统一模型入口

  22. class Model:

  23.     def __init__(self, device, allowedList=None, logger=None, requestId=None, modeType=None):

  24.         try:

  25.             logger.info("########################加载{}########################, requestId:{}", modeType.value[2], requestId)

  26.             self.modeType = modeType

  27.             self.allowedList = allowedList

  28.             self.trtFlag_det = True,  # 检测模型是否采用TRT

  29.             self.trtFlag_seg = True,  # 分割模型是否采用TRT

  30.             if self.trtFlag_det and self.trtFlag_seg:

  31.                 gpu = get_all_gpu_ids()[int(device)]

  32.                 if '3090' in gpu.name:

  33.                     detweights = "../AIlib2/weights/%s/yolov5_3090_fp16.engine" % modeType.value[3]

  34.                     segweights = '../AIlib2/weights/%s/stdc_360X640_3090_fp16.engine' % modeType.value[3]

  35.                 elif '2080' in gpu.name:

  36.                     detweights = "../AIlib2/weights/%s/yolov5_2080Ti_fp16.engine" % modeType.value[3]

  37.                     segweights = '../AIlib2/weights/%s/stdc_360X640_2080Ti_fp16.engine' % modeType.value[3]

  38.                 elif '4090' in gpu.name:

  39.                     detweights = "../AIlib2/weights/%s/yolov5_4090_fp16.engine" % modeType.value[3]

  40.                     segweights = '../AIlib2/weights/%s/stdc_360X640_4090_fp16.engine' % modeType.value[3]

  41.                 elif 'A10' in gpu.name:

  42.                     detweights = "../AIlib2/weights/%s/yolov5_A10_fp16.engine" % modeType.value[3]

  43.                     segweights = '../AIlib2/weights/%s/stdc_360X640_A10_fp16.engine' % modeType.value[3]

  44.                 else:

  45.                     raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  46.             else:

  47.                 detweights = "../AIlib2/weights/conf/%s/yolov5.pt" % modeType.value[3]

  48.                 segweights = '../AIlib2/weights/conf/%s/stdc_360X640.pth' % modeType.value[3]

  49.             par = {

  50.                 'device': device,  # 显卡号,如果用TRT模型，只支持0（单显卡）

  51.                 'labelnames': "../AIlib2/weights/conf/%s/labelnames.json" % modeType.value[3],  # 检测类别对照表

  52.                 'trtFlag_det': self.trtFlag_det,  # 检测模型是否采用TRT

  53.                 'trtFlag_seg': self.trtFlag_seg,  # 分割模型是否采用TRT

  54.                 'Detweights': detweights,  # 检测模型路径

  55.                 'seg_nclass': 2,  # 分割模型类别数目，默认2类

  56.                 'segPar': {

  57.                     'modelSize': (640, 360),

  58.                     'mean': (0.485, 0.456, 0.406),

  59.                     'std' : (0.229, 0.224, 0.225),

  60.                     'numpy': False,

  61.                     'RGB_convert_first': True

  62.                 },  # 分割模型预处理参数

  63.                 'Segweights': segweights,  # 分割模型权重位置

  64.                 'postFile': '../AIlib2/weights/conf/%s/para.json' % modeType.value[3]  # 后处理参数文件

  65.             }

  66.             if ModelType.WATER_SURFACE_MODEL == modeType:

  67.                 par['segRegionCnt'] = 1

  68.                 par['slopeIndex'] = [5, 6, 7]

  69.             if ModelType.TRAFFIC_FARM_MODEL == modeType:

  70.                 par['segRegionCnt'] = 2

  71.                 par['slopeIndex'] = [5, 6, 7]

  72.             else:

  73.                 par['segRegionCnt'] = 0

  74.                 par['slopeIndex'] = []

  75.                 par['segPar'] = None

  76.                 par['Segweights'] = None

  77.             self.device = select_device(par.get('device'))

  78.             self.names = get_labelnames(par.get('labelnames'))

  79.             self.half = self.device.type != 'cpu'

  80.             Detweights = par.get('Detweights')  # 升级后的检测模型

  81.             if self.trtFlag_det:

  82.                 with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:

  83.                     self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象

  84.                     print('####load TRT model :%s'%(Detweights))

  85.             else:

  86.                 self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  87.                 if self.half:

  88.                     self.model.half()

  89.             seg_nclass = par.get('seg_nclass')

  90.             self.segPar = par.get('segPar')

  91.             Segweights = par.get('Segweights')

  92.             if Segweights is not None:

  93.                 if self.trtFlag_seg:

  94.                     log = trt.Logger(trt.Logger.ERROR)

  95.                     with open(Segweights, "rb") as f, trt.Runtime(log) as runtime:

  96.                         self.segmodel = runtime.deserialize_cuda_engine(f.read())  # 输入trt本地文件，返回ICudaEngine对象

  97.                     print('############locad seg model trt success#######')

  98.                 else:

  99.                     self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device)

  100.                     print('############locad seg model pth success#######')

  101.             else:

  102.                 self.segmodel = None

  103.             self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  104.             self.objectPar = {

  105.                 'half': self.half,

  106.                 'device': self.device,

  107.                 'conf_thres': self.conf_thres,

  108.                 'iou_thres': self.iou_thres,

  109.                 'allowedList': self.allowedList,

  110.                 'slopeIndex': par.get('slopeIndex'),

  111.                 'segRegionCnt': par.get('segRegionCnt'),

  112.                 'trtFlag_det': self.trtFlag_det,

  113.                 'trtFlag_seg': self.trtFlag_seg

  114.             }

  115.             self.label_arraylist = None

  116.             self.digitFont = None

  117.         except Exception as ee:

  118.             logger.exception("异常：{}, requestId:{}", ee, requestId)

  119.             raise ee

  120. 

  121. 

  122.     def process(self, frame, width=1920):

  123.         if self.label_arraylist is None:

  124.             fontsize = int(width / 1920 * 40)

  125.             line_thickness = 1

  126.             boxLine_thickness = 1

  127.             waterLineWidth = 1

  128.             if width >= 960:

  129.                 line_thickness = int(round(width / 1920 * 3) - 1)

  130.                 boxLine_thickness = int(round(width / 1920 * 3))

  131.                 waterLineWidth = int(round(width / 1920 * 3))

  132.             numFontSize = float(format(width / 1920 * 1.1, '.1f'))  # 文字大小

  133.             self.digitFont = {'line_thickness': line_thickness,

  134.                               'boxLine_thickness': boxLine_thickness,

  135.                               'fontSize': numFontSize,

  136.                               'waterLineColor': (0, 255, 255),

  137.                               'waterLineWidth': waterLineWidth}

  138.             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  139.                                                     fontpath="../AIlib2/conf/platech.ttf")

  140.         if ModelType.WATER_SURFACE_MODEL == self.modeType or ModelType.TRAFFIC_FARM_MODEL == self.modeType:

  141.             return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  142.                               self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar)

  143.         else:

  144.             return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  145.                                      self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,

  146.                                      self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)

  147. 

  148. # # 河道模型

  149. # class SZModel:

  150. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  151. #         try:

  152. #             logger.info("########################加载河道模型########################, requestId:{}", requestId)

  153. #             self.allowedList = allowedList

  154. #             trtFlag_det = False,  # 检测模型是否采用TRT

  155. #             trtFlag_seg = False,  # 分割模型是否采用TRT

  156. #             if trtFlag_det and trtFlag_seg:

  157. #                 gpu = get_all_gpu_ids()[int(device)]

  158. #                 if '3090' in gpu.name:

  159. #                     det_gpuname = '_3090_fp16.engine'

  160. #                     seg_gpuname = '_3090_fp16.engine'

  161. #                 elif '2080' in gpu.name:

  162. #                     det_gpuname = '_2080Ti_fp16.engine'

  163. #                     seg_gpuname = '_2080Ti_fp16.engine'

  164. #                 elif '4090' in gpu.name:

  165. #                     det_gpuname = '_4090_fp16.engine'

  166. #                     seg_gpuname = '_4090_fp16.engine'

  167. #                 elif 'A10' in gpu.name:

  168. #                     det_gpuname = '_A10_fp16.engine'

  169. #                     seg_gpuname = '_A10_fp16.engine'

  170. #                 else:

  171. #                     raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  172. #             else:

  173. #                 det_gpuname = '.pt'

  174. #                 seg_gpuname = '.pth'

  175. #             par = {

  176. #                 'device': device,  # 显卡号,如果用TRT模型，只支持0（单显卡）

  177. #                 'labelnames': "../AIlib2/weights/river/labelnames.json",  # 检测类别对照表

  178. #                 'trtFlag_det': trtFlag_det,  # 检测模型是否采用TRT

  179. #                 'trtFlag_seg': trtFlag_seg,  # 分割模型是否采用TRT

  180. #                 'Detweights': "../AIlib2/weights/river/yolov5%s" % det_gpuname,  # 检测模型路径

  181. #                 'slopeIndex': [5, 6, 7],  # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。

  182. #                 'seg_nclass': 2,  # 分割模型类别数目，默认2类

  183. #                 'segRegionCnt': 1,  # 分割模型结果需要保留的等值线数目

  184. #                 'segPar': {

  185. #                     'modelSize': (640,360),

  186. #                     'mean': (0.485, 0.456, 0.406),

  187. #                     'std' : (0.229, 0.224, 0.225),

  188. #                     'numpy': False,

  189. #                     'RGB_convert_first': True

  190. #                 },  # 分割模型预处理参数

  191. #                 'Segweights': '../AIlib2/weights/river/stdc_360X640%s' % seg_gpuname,  # 分割模型权重位置

  192. #                 'postFile': '../AIlib2/weights/river/para.json'  # 后处理参数文件

  193. #             }

  194. #             self.device = select_device(par.get('device'))

  195. #             self.names = get_labelnames(par.get('labelnames'))

  196. #             half = self.device.type != 'cpu'

  197. #             Detweights = par.get('Detweights')  # 升级后的检测模型

  198. #             if trtFlag_det:

  199. #                 log = trt.Logger(trt.Logger.ERROR)

  200. #                 with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:

  201. #                     self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象

  202. #                 print('############locad det model trt success#######')

  203. #             else:

  204. #                 self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  205. #                 print('############locad det model pth success#######')

  206. #                 if half:

  207. #                     self.model.half()

  208. #             seg_nclass = par.get('seg_nclass')

  209. #             self.segPar = par.get('segPar')

  210. #             Segweights = par.get('Segweights')

  211. #             if trtFlag_seg:

  212. #                 log = trt.Logger(trt.Logger.ERROR)

  213. #                 with open(Segweights, "rb") as f, trt.Runtime(log) as runtime:

  214. #                     self.segmodel = runtime.deserialize_cuda_engine(f.read())  # 输入trt本地文件，返回ICudaEngine对象

  215. #                 print('############locad seg model trt success#######')

  216. #             else:

  217. #                 self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device)

  218. #                 print('############locad seg model pth success#######')

  219. #

  220. #             self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  221. #             self.objectPar = {

  222. #                 'half': half,

  223. #                 'device': self.device,

  224. #                 'conf_thres': self.conf_thres,

  225. #                 'iou_thres': self.iou_thres,

  226. #                 'allowedList': self.allowedList,

  227. #                 'slopeIndex': par.get('slopeIndex'),

  228. #                 'segRegionCnt': par.get('segRegionCnt'),

  229. #                 'trtFlag_det': trtFlag_det,

  230. #                 'trtFlag_seg': trtFlag_seg

  231. #             }

  232. #             self.label_arraylist = None

  233. #             self.digitFont = None

  234. #         except Exception as ee:

  235. #             logger.exception("异常：{}, requestId:{}", ee, requestId)

  236. #

  237. #

  238. #

  239. #     # names, label_arraylist, rainbows, conf_thres, iou_thres

  240. #     def process(self, frame, width=1920):

  241. #         if self.label_arraylist is None:

  242. #             fontsize = int(width / 1920 * 40)

  243. #             line_thickness = 1

  244. #             boxLine_thickness = 1

  245. #             waterLineWidth = 1

  246. #             if width >= 960:

  247. #                 line_thickness = int(round(width / 1920 * 3) - 1)

  248. #                 boxLine_thickness = int(round(width / 1920 * 3))

  249. #                 waterLineWidth = int(round(width / 1920 * 3))

  250. #             numFontSize = float(format(width / 1920 * 1.1, '.1f'))  # 文字大小

  251. #             self.digitFont = {'line_thickness': line_thickness,

  252. #                               'boxLine_thickness': boxLine_thickness,

  253. #                               'fontSize': numFontSize,

  254. #                               'waterLineColor': (0, 255, 255),

  255. #                               'waterLineWidth': waterLineWidth}

  256. #             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  257. #                                                     fontpath="../AIlib2/conf/platech.ttf")

  258. #         return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  259. #                           self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar)

  260. #

  261. #

  262. # # 森林模型

  263. # class LCModel:

  264. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  265. #         logger.info("########################加载森林模型########################, requestId:{}", requestId)

  266. #         self.allowedList = allowedList

  267. #         self.trtFlag_det = True,  # 检测模型是否采用TRT

  268. #         if self.trtFlag_det:

  269. #             gpu = get_all_gpu_ids()[int(device)]

  270. #             if '3090' in gpu.name:

  271. #                 det_gpuname = '_3090_fp16.engine'

  272. #             elif '2080' in gpu.name:

  273. #                 det_gpuname = '_2080Ti_fp16.engine'

  274. #             elif '4090' in gpu.name:

  275. #                 det_gpuname = '_4090_fp16.engine'

  276. #             elif 'A10' in gpu.name:

  277. #                 det_gpuname = '_A10_fp16.engine'

  278. #             else:

  279. #                 raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  280. #         else:

  281. #             det_gpuname = '.pt'

  282. #         par = {

  283. #             'device': device,   # 显卡号,如果用TRT模型，只支持0（单显卡）

  284. #             'labelnames': "../AIlib2/weights/forest/labelnames.json",  # 检测类别对照表

  285. #             'Detweights': "../AIlib2/weights/forest/yolov5%s" % det_gpuname,  # 检测模型路径

  286. #             'slopeIndex': [],  # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。

  287. #             'seg_nclass': 2,  # 分割模型类别数目，默认2类

  288. #             'segRegionCnt': 0,  # 分割模型结果需要保留的等值线数目

  289. #             'segPar': None,  # 分割模型预处理参数

  290. #             'Segweights': None,  # 分割模型权重位置

  291. #             'postFile': '../AIlib2/weights/forest/para.json'  # 后处理参数文件

  292. #         }

  293. #         self.device = select_device(par.get('device'))

  294. #         self.half = self.device.type != 'cpu'  # half precision only supported on CUDA

  295. #         Detweights = par.get('Detweights')

  296. #         if self.trtFlag_det:

  297. #             log = trt.Logger(trt.Logger.ERROR)

  298. #             with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:

  299. #                 self.model=runtime.deserialize_cuda_engine(f.read())

  300. #             print('####load TRT model :%s'%(Detweights))

  301. #         else:

  302. #             self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  303. #             if self.half:

  304. #                 self.model.half()

  305. #         self.names = get_labelnames(par.get('labelnames'))

  306. #         self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  307. #         self.digitFont = None

  308. #         self.label_arraylist = None

  309. #         self.segmodel = None

  310. #

  311. #         # names, label_arraylist, rainbows, conf_thres, iou_thres

  312. #     def process(self, frame, width=1920):

  313. #         if self.label_arraylist is None:

  314. #             fontsize = int(width / 1920 * 40)

  315. #             line_thickness = 1

  316. #             boxLine_thickness = 1

  317. #             waterLineWidth = 1

  318. #             if width >= 960:

  319. #                 line_thickness = int(round(width / 1920 * 3) - 1)

  320. #                 boxLine_thickness = int(round(width / 1920 * 3))

  321. #                 waterLineWidth = int(round(width / 1920 * 3))

  322. #             numFontSize = float(format(width / 1920 * 1.1, '.1f'))  # 文字大小

  323. #             self.digitFont = {'line_thickness': line_thickness,

  324. #                               'boxLine_thickness': boxLine_thickness,

  325. #                               'fontSize': numFontSize,

  326. #                               'waterLineColor': (0, 255, 255),

  327. #                               'waterLineWidth': waterLineWidth}

  328. #             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  329. #                                                     fontpath="../AIlib2/conf/platech.ttf")

  330. #         return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  331. #                                  self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,

  332. #                                  self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)

  333. #

  334. #

  335. # # 车辆模型

  336. # class VehicleModel:

  337. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  338. #         logger.info("########################加载车辆模型########################, requestId:{}", requestId)

  339. #         self.allowedList = allowedList

  340. #         self.trtFlag_det = True,  # 检测模型是否采用TRT

  341. #         if self.trtFlag_det:

  342. #             gpu = get_all_gpu_ids()[int(device)]

  343. #             if '3090' in gpu.name:

  344. #                 det_gpuname = '_3090_fp16.engine'

  345. #             elif '2080' in gpu.name:

  346. #                 det_gpuname = '_2080Ti_fp16.engine'

  347. #             elif '4090' in gpu.name:

  348. #                 det_gpuname = '_4090_fp16.engine'

  349. #             elif 'A10' in gpu.name:

  350. #                 det_gpuname = '_A10_fp16.engine'

  351. #             else:

  352. #                 raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  353. #         else:

  354. #             det_gpuname = '.pt'

  355. #         par = {

  356. #             'device': device,   # 显卡号,如果用TRT模型，只支持0（单显卡）

  357. #             'labelnames': "../AIlib2/weights/vehicle/labelnames.json",  # 检测类别对照表

  358. #             'Detweights': "../AIlib2/weights/vehicle/yolov5%s" % det_gpuname,  # 检测模型路径

  359. #             'slopeIndex': [],  # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。

  360. #             'seg_nclass': 2,  # 分割模型类别数目，默认2类

  361. #             'segRegionCnt': 0,  # 分割模型结果需要保留的等值线数目

  362. #             'segPar': None,  # 分割模型预处理参数

  363. #             'Segweights': None,  # 分割模型权重位置

  364. #             'postFile': '../AIlib2/weights/vehicle/para.json'  # 后处理参数文件

  365. #         }

  366. #         self.device = select_device(par.get('device'))

  367. #         self.half = self.device.type != 'cpu'  # half precision only supported on CUDA

  368. #         Detweights = par.get('Detweights')

  369. #         if self.trtFlag_det:

  370. #             log = trt.Logger(trt.Logger.ERROR)

  371. #             with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:

  372. #                 self.model=runtime.deserialize_cuda_engine(f.read())

  373. #             print('####load TRT model :%s'%(Detweights))

  374. #         else:

  375. #             self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  376. #             if self.half:

  377. #                 self.model.half()

  378. #         self.names = get_labelnames(par.get('labelnames'))

  379. #         self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  380. #         self.digitFont = None

  381. #         self.label_arraylist = None

  382. #         self.segmodel = None

  383. #

  384. #         # names, label_arraylist, rainbows, conf_thres, iou_thres

  385. #     def process(self, frame, width=1920):

  386. #         if self.label_arraylist is None:

  387. #             fontsize = int(width / 1920 * 40)

  388. #             line_thickness = 1

  389. #             boxLine_thickness = 1

  390. #             waterLineWidth = 1

  391. #             if width >= 960:

  392. #                 line_thickness = int(round(width / 1920 * 3) - 1)

  393. #                 boxLine_thickness = int(round(width / 1920 * 3))

  394. #                 waterLineWidth = int(round(width / 1920 * 3))

  395. #             numFontSize = float(format(width / 1920 * 1.1, '.1f'))  # 文字大小

  396. #             self.digitFont = {'line_thickness': line_thickness,

  397. #                               'boxLine_thickness': boxLine_thickness,

  398. #                               'fontSize': numFontSize,

  399. #                               'waterLineColor': (0, 255, 255),

  400. #                               'waterLineWidth': waterLineWidth}

  401. #             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  402. #                                                     fontpath="../AIlib2/conf/platech.ttf")

  403. #         return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  404. #                                  self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,

  405. #                                  self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)

  406. 

  407. 

  408. # # 行人模型

  409. # class PedestrianModel:

  410. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  411. #         logger.info("########################加载行人模型########################, requestId:{}", requestId)

  412. #         self.allowedList = allowedList

  413. #         self.trtFlag_det = True,  # 检测模型是否采用TRT

  414. #         if self.trtFlag_det:

  415. #             gpu = get_all_gpu_ids()[int(device)]

  416. #             if '3090' in gpu.name:

  417. #                 det_gpuname = '_3090_fp16.engine'

  418. #             elif '2080' in gpu.name:

  419. #                 det_gpuname = '_2080Ti_fp16.engine'

  420. #             elif '4090' in gpu.name:

  421. #                 det_gpuname = '_4090_fp16.engine'

  422. #             elif 'A10' in gpu.name:

  423. #                 det_gpuname = '_A10_fp16.engine'

  424. #             else:

  425. #                 raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  426. #         else:

  427. #             det_gpuname = '.pt'

  428. #         par = {

  429. #             'device': device,   # 显卡号,如果用TRT模型，只支持0（单显卡）

  430. #             'labelnames': "../AIlib2/weights/pedestrian/labelnames.json",  # 检测类别对照表

  431. #             'Detweights': "../AIlib2/weights/pedestrian/yolov5%s" % det_gpuname,  # 检测模型路径

  432. #             'slopeIndex': [],  # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。

  433. #             'seg_nclass': 2,  # 分割模型类别数目，默认2类

  434. #             'segRegionCnt': 0,  # 分割模型结果需要保留的等值线数目

  435. #             'segPar': None,  # 分割模型预处理参数

  436. #             'Segweights': None,  # 分割模型权重位置

  437. #             'postFile': '../AIlib2/weights/pedestrian/para.json'  # 后处理参数文件

  438. #         }

  439. #         self.device = select_device(par.get('device'))

  440. #         self.half = self.device.type != 'cpu'  # half precision only supported on CUDA

  441. #         Detweights = par.get('Detweights')

  442. #         if self.trtFlag_det:

  443. #             log = trt.Logger(trt.Logger.ERROR)

  444. #             with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:

  445. #                 self.model=runtime.deserialize_cuda_engine(f.read())

  446. #             print('####load TRT model :%s'%(Detweights))

  447. #         else:

  448. #             self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  449. #             if self.half:

  450. #                 self.model.half()

  451. #         self.names = get_labelnames(par.get('labelnames'))

  452. #         self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  453. #         self.digitFont = None

  454. #         self.label_arraylist = None

  455. #         self.segmodel = None

  456. #

  457. #         # names, label_arraylist, rainbows, conf_thres, iou_thres

  458. #     def process(self, frame, width=1920):

  459. #         if self.label_arraylist is None:

  460. #             fontsize = int(width / 1920 * 40)

  461. #             line_thickness = 1

  462. #             boxLine_thickness = 1

  463. #             waterLineWidth = 1

  464. #             if width >= 960:

  465. #                 line_thickness = int(round(width / 1920 * 3) - 1)

  466. #                 boxLine_thickness = int(round(width / 1920 * 3))

  467. #                 waterLineWidth = int(round(width / 1920 * 3))

  468. #             numFontSize = float(format(width / 1920 * 1.1, '.1f'))  # 文字大小

  469. #             self.digitFont = {'line_thickness': line_thickness,

  470. #                               'boxLine_thickness': boxLine_thickness,

  471. #                               'fontSize': numFontSize,

  472. #                               'waterLineColor': (0, 255, 255),

  473. #                               'waterLineWidth': waterLineWidth}

  474. #             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  475. #                                                     fontpath="../AIlib2/conf/platech.ttf")

  476. #         return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  477. #                                  self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,

  478. #                                  self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)

  479. #

  480. #

  481. # # 烟火模型

  482. # class SmogfireModel:

  483. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  484. #         logger.info("########################加载烟火模型########################, requestId:{}", requestId)

  485. #         self.allowedList = allowedList

  486. #         self.trtFlag_det = True,  # 检测模型是否采用TRT

  487. #         if self.trtFlag_det:

  488. #             gpu = get_all_gpu_ids()[int(device)]

  489. #             if '3090' in gpu.name:

  490. #                 det_gpuname = '_3090_fp16.engine'

  491. #             elif '2080' in gpu.name:

  492. #                 det_gpuname = '_2080Ti_fp16.engine'

  493. #             elif '4090' in gpu.name:

  494. #                 det_gpuname = '_4090_fp16.engine'

  495. #             elif 'A10' in gpu.name:

  496. #                 det_gpuname = '_A10_fp16.engine'

  497. #             else:

  498. #                 raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  499. #         else:

  500. #             det_gpuname = '.pt'

  501. #         par = {

  502. #             'device': device,   # 显卡号,如果用TRT模型，只支持0（单显卡）

  503. #             'labelnames': "../AIlib2/weights/smogfire/labelnames.json",  # 检测类别对照表

  504. #             'Detweights': "../AIlib2/weights/smogfire/yolov5%s" % det_gpuname,  # 检测模型路径

  505. #             'slopeIndex': [],  # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。

  506. #             'seg_nclass': 2,  # 分割模型类别数目，默认2类

  507. #             'segRegionCnt': 0,  # 分割模型结果需要保留的等值线数目

  508. #             'segPar': None,  # 分割模型预处理参数

  509. #             'Segweights': None,  # 分割模型权重位置

  510. #             'postFile': '../AIlib2/weights/smogfire/para.json'  # 后处理参数文件

  511. #         }

  512. #         self.device = select_device(par.get('device'))

  513. #         self.half = self.device.type != 'cpu'  # half precision only supported on CUDA

  514. #         Detweights = par.get('Detweights')

  515. #         if self.trtFlag_det:

  516. #             log = trt.Logger(trt.Logger.ERROR)

  517. #             with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:

  518. #                 self.model=runtime.deserialize_cuda_engine(f.read())

  519. #             print('####load TRT model :%s'%(Detweights))

  520. #         else:

  521. #             self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  522. #             if self.half:

  523. #                 self.model.half()

  524. #         self.names = get_labelnames(par.get('labelnames'))

  525. #         self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  526. #         self.digitFont = None

  527. #         self.label_arraylist = None

  528. #         self.segmodel = None

  529. #

  530. #         # names, label_arraylist, rainbows, conf_thres, iou_thres

  531. #     def process(self, frame, width=1920):

  532. #         if self.label_arraylist is None:

  533. #             fontsize = int(width / 1920 * 40)

  534. #             line_thickness = 1

  535. #             boxLine_thickness = 1

  536. #             waterLineWidth = 1

  537. #             if width >= 960:

  538. #                 line_thickness = int(round(width / 1920 * 3) - 1)

  539. #                 boxLine_thickness = int(round(width / 1920 * 3))

  540. #                 waterLineWidth = int(round(width / 1920 * 3))

  541. #             numFontSize = float(format(width / 1920 * 1.1, '.1f'))  # 文字大小

  542. #             self.digitFont = {'line_thickness': line_thickness,

  543. #                               'boxLine_thickness': boxLine_thickness,

  544. #                               'fontSize': numFontSize,

  545. #                               'waterLineColor': (0, 255, 255),

  546. #                               'waterLineWidth': waterLineWidth}

  547. #             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  548. #                                                     fontpath="../AIlib2/conf/platech.ttf")

  549. #         return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  550. #                                  self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres,

  551. #                                  self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det)

  552. #

  553. # # 交通模型

  554. # class RFModel:

  555. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  556. #         logger.info("########################加载交通模型########################, requestId:{}", requestId)

  557. #         self.allowedList = allowedList

  558. #         trtFlag_det = True,  # 检测模型是否采用TRT

  559. #         trtFlag_seg = True,  # 分割模型是否采用TRT

  560. #         if trtFlag_det and trtFlag_seg:

  561. #             gpu = get_all_gpu_ids()[int(device)]

  562. #             if '3090' in gpu.name:

  563. #                 det_gpuname = '_3090_fp16.engine'

  564. #                 seg_gpuname = '_3090_fp16.engine'

  565. #             elif '2080' in gpu.name:

  566. #                 det_gpuname = '_2080Ti_fp16.engine'

  567. #                 seg_gpuname = '_2080Ti_fp16.engine'

  568. #             elif '4090' in gpu.name:

  569. #                 det_gpuname = '_4090_fp16.engine'

  570. #                 seg_gpuname = '_4090_fp16.engine'

  571. #             elif 'A10' in gpu.name:

  572. #                 det_gpuname = '_A10_fp16.engine'

  573. #                 seg_gpuname = '_A10_fp16.engine'

  574. #             else:

  575. #                 raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name)

  576. #         else:

  577. #             det_gpuname = '.pt'

  578. #             seg_gpuname = '.pth'

  579. #         par = {

  580. #             'device': device,  # 显卡号,如果用TRT模型，只支持0（单显卡）

  581. #             'labelnames': "../AIlib2/weights/road/labelnames.json",  # 检测类别对照表

  582. #             'trtFlag_det': trtFlag_det,  # 检测模型是否采用TRT

  583. #             'trtFlag_seg': trtFlag_seg,  # 分割模型是否采用TRT

  584. #             'Detweights': "../AIlib2/weights/road/yolov5%s" % det_gpuname,  # 检测模型路径

  585. #             'slopeIndex': [],  # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。

  586. #             'seg_nclass': 2,  # 分割模型类别数目，默认2类

  587. #             'segRegionCnt': 2,  # 分割模型结果需要保留的等值线数目

  588. #             'segPar': {

  589. #                 'modelSize': (640,360),

  590. #                 'mean': (0.485, 0.456, 0.406),

  591. #                 'std' : (0.229, 0.224, 0.225),

  592. #                 'numpy': False,

  593. #                 'RGB_convert_first': True

  594. #             },  # 分割模型预处理参数

  595. #             'Segweights': '../AIlib2/weights/road/stdc_360X640%s' % seg_gpuname,  # 分割模型权重位置

  596. #             'postFile': '../AIlib2/weights/road/para.json'  # 后处理参数文件

  597. #         }

  598. #         self.device = select_device(par.get('device'))

  599. #         self.names = get_labelnames(par.get('labelnames'))

  600. #         half = self.device.type != 'cpu'

  601. #         Detweights = par.get('Detweights')  # 升级后的检测模型

  602. #         if trtFlag_det:

  603. #             log = trt.Logger(trt.Logger.ERROR)

  604. #             with open(Detweights, "rb") as f, trt.Runtime(log) as runtime:

  605. #                 self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象

  606. #             print('############locad det model trt success#######')

  607. #         else:

  608. #             self.model = attempt_load(Detweights, map_location=self.device)  # load FP32 model

  609. #             print('############locad det model pth success#######')

  610. #             if half:

  611. #                 self.model.half()

  612. #         seg_nclass = par.get('seg_nclass')

  613. #         self.segPar = par.get('segPar')

  614. #         Segweights = par.get('Segweights')

  615. #         if trtFlag_seg:

  616. #             log = trt.Logger(trt.Logger.ERROR)

  617. #             with open(Segweights, "rb") as f, trt.Runtime(log) as runtime:

  618. #                 self.segmodel = runtime.deserialize_cuda_engine(f.read())  # 输入trt本地文件，返回ICudaEngine对象

  619. #             print('############locad seg model trt success#######')

  620. #         else:

  621. #             self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device)

  622. #             print('############locad seg model pth success#######')

  623. #

  624. #         self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile'))

  625. #         self.objectPar = {

  626. #             'half': half,

  627. #             'device': self.device,

  628. #             'conf_thres': self.conf_thres,

  629. #             'iou_thres': self.iou_thres,

  630. #             'allowedList': self.allowedList,

  631. #             'slopeIndex': par.get('slopeIndex'),

  632. #             'segRegionCnt': par.get('segRegionCnt'),

  633. #             'trtFlag_det': trtFlag_det,

  634. #             'trtFlag_seg': trtFlag_seg

  635. #         }

  636. #         self.label_arraylist = None

  637. #         self.digitFont = None

  638. #

  639. #     def process(self, frame, width=1920):

  640. #         if self.label_arraylist is None:

  641. #             fontsize = int(width / 1920 * 40)

  642. #             line_thickness = 1

  643. #             boxLine_thickness = 1

  644. #             waterLineWidth = 1

  645. #             if width >= 960:

  646. #                 line_thickness = int(round(width / 1920 * 3) - 1)

  647. #                 boxLine_thickness = int(round(width / 1920 * 2))

  648. #                 waterLineWidth = int(round(width / 1920 * 3))

  649. #             numFontSize = float(format(width / 1920 * 1.1, '.1f'))

  650. #

  651. #             self.digitFont = {'line_thickness': line_thickness,

  652. #                               'boxLine_thickness': boxLine_thickness,

  653. #                               'fontSize': numFontSize,

  654. #                               'waterLineColor': (0, 255, 255),

  655. #                               'waterLineWidth': waterLineWidth}

  656. #             self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize,

  657. #                                                     fontpath="../AIlib2/conf/platech.ttf")

  658. #         return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist,

  659. #                           self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar)

  660. 

  661. # class IMModel:

  662. #     def __init__(self, device, allowedList=None, logger=None, requestId=None):

  663. #         logger.info("########################加载防疫模型########################, requestId:{}", requestId)

  664. #         self.allowedList = allowedList

  665. #         self.img_type = 'code'

  666. #         self.par = {'code': {'weights': '../AIlib2/weights/jkm/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10},

  667. #                     'plate': {'weights': '../AIlib2/weights/jkm/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1},

  668. #                     'conf_thres': 0.4,

  669. #                     'iou_thres': 0.45,

  670. #                     'device': 'cuda:%s' % device,

  671. #                     'plate_dilate': (0.5, 0.3)

  672. #                     }

  673. #         self.device = torch.device(self.par['device'])

  674. #         self.model = torch.jit.load(self.par[self.img_type]['weights'])

  675. #

  676. #     def process(self, frame):

  677. #         img, padInfos = pre_process(frame, self.device)  ##预处理

  678. #         pred = self.model(img)  ##模型推理

  679. #         boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'],

  680. #                              iou_thres=self.par['iou_thres'], nc=self.par[self.img_type]['nc'])  # 后处理

  681. #         dataBack = get_return_data(frame, boxes, modelType=self.img_type, plate_dilate=self.par['plate_dilate'])

  682. #         return dataBack

  683. 

  684. 

  685. # 车牌分割模型、健康码、行程码分割模型

  686. class IMModel:

  687. 

  688.     def __init__(self, device, allowedList=None, logger=None, requestId=None, modeType=None):

  689.         logger.info("########################加载{}########################, requestId:{}", modeType.value[2], requestId)

  690.         self.allowedList = allowedList

  691.         if ModelType.EPIDEMIC_PREVENTION_MODEL == modeType:

  692.             self.img_type = 'code'  ## code,plate

  693.         if ModelType.PLATE_MODEL == modeType:

  694.             self.img_type = 'plate'  ## code,plate

  695.         self.par = {'code': {'weights': '../AIlib2/weights/conf/jkm/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10},

  696.                     'plate': {'weights': '../AIlib2/weights/conf/jkm/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1},

  697.                     'conf_thres': 0.4,

  698.                     'iou_thres': 0.45,

  699.                     'device': 'cuda:%s' % device,

  700.                     'plate_dilate': (0.5, 0.3)

  701.                     }

  702.         self.device = torch.device(self.par['device'])

  703.         self.model = torch.jit.load(self.par[self.img_type]['weights'])

  704. 

  705.     def process(self, frame):

  706.         img, padInfos = pre_process(frame, self.device)  ##预处理

  707.         pred = self.model(img)  ##模型推理

  708.         boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'],

  709.                              iou_thres=self.par['iou_thres'], nc=self.par[self.img_type]['nc'])  # 后处理

  710.         dataBack = get_return_data(frame, boxes, modelType=self.img_type, plate_dilate=self.par['plate_dilate'])

  711.         return dataBack