|
|
|
@@ -1,72 +1,104 @@ |
|
|
|
# -*- coding: utf-8 -*- |
|
|
|
import sys |
|
|
|
|
|
|
|
import torch |
|
|
|
|
|
|
|
from util.GPUtils import get_all_gpu_ids |
|
|
|
|
|
|
|
sys.path.extend(['..', '../AIlib']) |
|
|
|
sys.path.extend(['..', '../AIlib2']) |
|
|
|
from AI import AI_process,AI_process_forest,get_postProcess_para |
|
|
|
import cv2,os,time |
|
|
|
from segutils.segmodel import SegModel |
|
|
|
from segutils.segmodel_trt import SegModel_STDC_trt |
|
|
|
from segutils.trtUtils import DetectMultiBackend |
|
|
|
from utils.torch_utils import select_device |
|
|
|
from models.experimental import attempt_load |
|
|
|
from utilsK.queRiver import get_labelnames, get_label_arrays, post_process_ |
|
|
|
from AI import AI_process, AI_process_v2, AI_process_forest, get_postProcess_para |
|
|
|
from utilsK.jkmUtils import pre_process, post_process, get_return_data |
|
|
|
from utils.torch_utils import select_device |
|
|
|
from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images |
|
|
|
import numpy as np |
|
|
|
import torch,glob |
|
|
|
import tensorrt as trt |
|
|
|
from utilsK.masterUtils import get_needed_objectsIndex |
|
|
|
from utilsK.jkmUtils import pre_process,post_process,get_return_data |
|
|
|
|
|
|
|
|
|
|
|
class Model(): |
|
|
|
def __init__(self, device, allowedList=None): |
|
|
|
##预先设置的参数 |
|
|
|
self.device_ = device ##选定模型，可选 cpu,'0','1' |
|
|
|
# 河道模型 |
|
|
|
class SZModel: |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载河道模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
|
|
|
|
|
|
|
|
# 水面模型 |
|
|
|
class SZModel(Model): |
|
|
|
def __init__(self, device, allowedList=None): |
|
|
|
super().__init__(device, allowedList) |
|
|
|
postFile = '../AIlib/conf/para.json' |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(postFile) |
|
|
|
self.label_arraylist = None |
|
|
|
self.digitFont = None |
|
|
|
labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表 |
|
|
|
self.names = get_labelnames(labelnames) |
|
|
|
self.trtFlag_det = True |
|
|
|
self.trtFlag_seg = True |
|
|
|
self.device = select_device(self.device_) |
|
|
|
self.half = self.device.type != 'cpu' |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
gpuname = '3090' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
gpuname = '2080Ti' |
|
|
|
trtFlag_det = True, # 检测模型是否采用TRT |
|
|
|
trtFlag_seg = True, # 分割模型是否采用TRT |
|
|
|
if trtFlag_det and trtFlag_seg: |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
det_gpuname = '_3090_fp16.engine' |
|
|
|
seg_gpuname = '_3090_fp16.engine' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
det_gpuname = '_2080Ti_fp16.engine' |
|
|
|
seg_gpuname = '_2080Ti_fp16.engine' |
|
|
|
elif 'A10' in gpu.name: |
|
|
|
det_gpuname = '_A10_fp16.engine' |
|
|
|
seg_gpuname = '_A10_fp16.engine' |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
if self.trtFlag_det: |
|
|
|
logger = trt.Logger(trt.Logger.ERROR) |
|
|
|
|
|
|
|
with open("../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine" % gpuname, "rb") as f, trt.Runtime( |
|
|
|
logger) as runtime: |
|
|
|
self.model = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
det_gpuname = '.pt' |
|
|
|
seg_gpuname = '.pth' |
|
|
|
par = { |
|
|
|
'device': device, # 显卡号,如果用TRT模型，只支持0（单显卡） |
|
|
|
'labelnames': "../AIlib2/weights/river/labelnames.json", # 检测类别对照表 |
|
|
|
'trtFlag_det': trtFlag_det, # 检测模型是否采用TRT |
|
|
|
'trtFlag_seg': trtFlag_seg, # 分割模型是否采用TRT |
|
|
|
'Detweights': "../AIlib2/weights/river/yolov5%s" % det_gpuname, # 检测模型路径 |
|
|
|
'slopeIndex': [5, 6, 7], # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。 |
|
|
|
'seg_nclass': 2, # 分割模型类别数目，默认2类 |
|
|
|
'segRegionCnt': 1, # 分割模型结果需要保留的等值线数目 |
|
|
|
'segPar': { |
|
|
|
'modelSize': (640,360), |
|
|
|
'mean': (0.485, 0.456, 0.406), |
|
|
|
'std' : (0.229, 0.224, 0.225), |
|
|
|
'numpy': False, |
|
|
|
'RGB_convert_first': True |
|
|
|
}, # 分割模型预处理参数 |
|
|
|
'Segweights': '../AIlib2/weights/river/stdc_360X640%s' % seg_gpuname, # 分割模型权重位置 |
|
|
|
'postFile': '../AIlib2/weights/river/para.json' # 后处理参数文件 |
|
|
|
} |
|
|
|
self.device = select_device(par.get('device')) |
|
|
|
self.names = get_labelnames(par.get('labelnames')) |
|
|
|
half = self.device.type != 'cpu' |
|
|
|
Detweights = par.get('Detweights') # 升级后的检测模型 |
|
|
|
if trtFlag_det: |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
print('############locad det model trt success#######') |
|
|
|
else: |
|
|
|
self.model = attempt_load("../AIlib/weights/yolov5/class8/bestcao.pt", |
|
|
|
map_location=self.device) # load FP32 model |
|
|
|
if self.half: |
|
|
|
self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model |
|
|
|
print('############locad det model pth success#######') |
|
|
|
if half: |
|
|
|
self.model.half() |
|
|
|
seg_nclass = 2 |
|
|
|
self.segPar = {'modelSize': (640, 360), 'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225), |
|
|
|
'numpy': False, 'RGB_convert_first': True} |
|
|
|
if self.trtFlag_seg: |
|
|
|
Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640_%s_fp16.engine' % gpuname ##升级的分割模型 |
|
|
|
logger = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Segweights, "rb") as f, trt.Runtime(logger) as runtime: |
|
|
|
seg_nclass = par.get('seg_nclass') |
|
|
|
self.segPar = par.get('segPar') |
|
|
|
Segweights = par.get('Segweights') |
|
|
|
if trtFlag_seg: |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Segweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.segmodel = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
print('############locad seg model trt success#######') |
|
|
|
else: |
|
|
|
Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' |
|
|
|
self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device) |
|
|
|
print('############locad seg model pth success#######') |
|
|
|
|
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile')) |
|
|
|
self.objectPar = { |
|
|
|
'half': half, |
|
|
|
'device': self.device, |
|
|
|
'conf_thres': self.conf_thres, |
|
|
|
'iou_thres': self.iou_thres, |
|
|
|
'allowedList': self.allowedList, |
|
|
|
'slopeIndex': par.get('slopeIndex'), |
|
|
|
'segRegionCnt': par.get('segRegionCnt'), |
|
|
|
'trtFlag_det': trtFlag_det, |
|
|
|
'trtFlag_seg': trtFlag_seg |
|
|
|
} |
|
|
|
self.label_arraylist = None |
|
|
|
self.digitFont = None |
|
|
|
|
|
|
|
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres |
|
|
|
def process(self, frame, width=1920): |
|
|
|
@@ -86,45 +118,56 @@ class SZModel(Model): |
|
|
|
'waterLineColor': (0, 255, 255), |
|
|
|
'waterLineWidth': waterLineWidth} |
|
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize, |
|
|
|
fontpath="../AIlib/conf/platech.ttf") |
|
|
|
fontpath="../AIlib2/conf/platech.ttf") |
|
|
|
return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist, |
|
|
|
self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres, |
|
|
|
self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det, |
|
|
|
trtFlag_seg=self.trtFlag_seg, segPar=self.segPar) |
|
|
|
self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar) |
|
|
|
|
|
|
|
|
|
|
|
# 森林模型 |
|
|
|
class LCModel(Model): |
|
|
|
def __init__(self, device, allowedList=None): |
|
|
|
super().__init__(device, allowedList) |
|
|
|
self.label_arraylist = None |
|
|
|
labelnames = "../AIlib/weights/forest/labelnames.json" |
|
|
|
self.names = get_labelnames(labelnames) |
|
|
|
postFile = '../AIlib/weights/forest/para.json' |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(postFile) |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
gpuname = '3090' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
gpuname = '2080Ti' |
|
|
|
class LCModel: |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载森林模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
self.trtFlag_det = True, # 检测模型是否采用TRT |
|
|
|
if self.trtFlag_det: |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
det_gpuname = '_3090_fp16.engine' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
det_gpuname = '_2080Ti_fp16.engine' |
|
|
|
elif 'A10' in gpu.name: |
|
|
|
det_gpuname = '_A10_fp16.engine' |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
self.device = select_device(self.device_) |
|
|
|
self.half = self.device.type != 'cpu' # half precision only supported on CUDA |
|
|
|
|
|
|
|
self.trtFlag_det = True # 是否采用TRT模型加速 |
|
|
|
det_gpuname = '.pt' |
|
|
|
par = { |
|
|
|
'device': device, # 显卡号,如果用TRT模型，只支持0（单显卡） |
|
|
|
'labelnames': "../AIlib2/weights/forest/labelnames.json", # 检测类别对照表 |
|
|
|
'Detweights': "../AIlib2/weights/forest/yolov5%s" % det_gpuname, # 检测模型路径 |
|
|
|
'slopeIndex': [], # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。 |
|
|
|
'seg_nclass': 2, # 分割模型类别数目，默认2类 |
|
|
|
'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目 |
|
|
|
'segPar': None, # 分割模型预处理参数 |
|
|
|
'Segweights': None, # 分割模型权重位置 |
|
|
|
'postFile': '../AIlib2/weights/forest/para.json' # 后处理参数文件 |
|
|
|
} |
|
|
|
self.device = select_device(par.get('device')) |
|
|
|
self.half = device.type != 'cpu' # half precision only supported on CUDA |
|
|
|
Detweights = par.get('Detweights') |
|
|
|
if self.trtFlag_det: |
|
|
|
detweights = "../AIlib/weights/forest/best_%s_fp16.engine" % gpuname |
|
|
|
logger = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(detweights, "rb") as f, trt.Runtime(logger) as runtime: |
|
|
|
self.model = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
print("1111111111111111111111111111111111111111111111111111") |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.model=runtime.deserialize_cuda_engine(f.read()) |
|
|
|
print('####load TRT model :%s'%(Detweights)) |
|
|
|
else: |
|
|
|
detweights = "../AIlib/weights/forest/best.pt" |
|
|
|
self.model = attempt_load(detweights, map_location=self.device) # load FP32 model |
|
|
|
self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model |
|
|
|
if self.half: |
|
|
|
self.model.half() |
|
|
|
print("222222222222222222222222222222222222222222222222222222222") |
|
|
|
self.names = get_labelnames(par.get('labelnames')) |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile')) |
|
|
|
self.digitFont = None |
|
|
|
self.label_arraylist = None |
|
|
|
self.segmodel = None |
|
|
|
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres |
|
|
|
@@ -145,47 +188,131 @@ class LCModel(Model): |
|
|
|
'waterLineColor': (0, 255, 255), |
|
|
|
'waterLineWidth': waterLineWidth} |
|
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize, |
|
|
|
fontpath="../AIlib/conf/platech.ttf") |
|
|
|
fontpath="../AIlib2/conf/platech.ttf") |
|
|
|
return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist, |
|
|
|
self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres, |
|
|
|
self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det) |
|
|
|
|
|
|
|
|
|
|
|
# 交通模型 |
|
|
|
class RFModel(Model): |
|
|
|
def __init__(self, device, allowedList=None): |
|
|
|
super().__init__(device, allowedList) |
|
|
|
self.label_arraylist = None |
|
|
|
labelnames = "../AIlib/weights/road/labelnames.json" |
|
|
|
self.names = get_labelnames(labelnames) |
|
|
|
postFile = '../AIlib/weights/road/para.json' |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(postFile) |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
gpuname = '3090' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
gpuname = '2080Ti' |
|
|
|
# 车辆模型 |
|
|
|
class VehicleModel: |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载车辆模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
self.trtFlag_det = True, # 检测模型是否采用TRT |
|
|
|
if self.trtFlag_det: |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
det_gpuname = '_3090_fp16.engine' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
det_gpuname = '_2080Ti_fp16.engine' |
|
|
|
elif 'A10' in gpu.name: |
|
|
|
det_gpuname = '_A10_fp16.engine' |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
else: |
|
|
|
det_gpuname = '.pt' |
|
|
|
par = { |
|
|
|
'device': device, # 显卡号,如果用TRT模型，只支持0（单显卡） |
|
|
|
'labelnames': "../AIlib2/weights/vehicle/labelnames.json", # 检测类别对照表 |
|
|
|
'Detweights': "../AIlib2/weights/vehicle/yolov5%s" % det_gpuname, # 检测模型路径 |
|
|
|
'slopeIndex': [], # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。 |
|
|
|
'seg_nclass': 2, # 分割模型类别数目，默认2类 |
|
|
|
'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目 |
|
|
|
'segPar': None, # 分割模型预处理参数 |
|
|
|
'Segweights': None, # 分割模型权重位置 |
|
|
|
'postFile': '../AIlib2/weights/vehicle/para.json' # 后处理参数文件 |
|
|
|
} |
|
|
|
self.device = select_device(par.get('device')) |
|
|
|
self.half = device.type != 'cpu' # half precision only supported on CUDA |
|
|
|
Detweights = par.get('Detweights') |
|
|
|
if self.trtFlag_det: |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.model=runtime.deserialize_cuda_engine(f.read()) |
|
|
|
print('####load TRT model :%s'%(Detweights)) |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model |
|
|
|
if self.half: |
|
|
|
self.model.half() |
|
|
|
self.names = get_labelnames(par.get('labelnames')) |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile')) |
|
|
|
self.digitFont = None |
|
|
|
self.label_arraylist = None |
|
|
|
self.segmodel = None |
|
|
|
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres |
|
|
|
def process(self, frame, width=1920): |
|
|
|
if self.label_arraylist is None: |
|
|
|
fontsize = int(width / 1920 * 40) |
|
|
|
line_thickness = 1 |
|
|
|
boxLine_thickness = 1 |
|
|
|
waterLineWidth = 1 |
|
|
|
if width >= 960: |
|
|
|
line_thickness = int(round(width / 1920 * 3) - 1) |
|
|
|
boxLine_thickness = int(round(width / 1920 * 3)) |
|
|
|
waterLineWidth = int(round(width / 1920 * 3)) |
|
|
|
numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小 |
|
|
|
self.digitFont = {'line_thickness': line_thickness, |
|
|
|
'boxLine_thickness': boxLine_thickness, |
|
|
|
'fontSize': numFontSize, |
|
|
|
'waterLineColor': (0, 255, 255), |
|
|
|
'waterLineWidth': waterLineWidth} |
|
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize, |
|
|
|
fontpath="../AIlib2/conf/platech.ttf") |
|
|
|
return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist, |
|
|
|
self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres, |
|
|
|
self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det) |
|
|
|
|
|
|
|
|
|
|
|
self.device = select_device(self.device_) |
|
|
|
self.half = self.device.type != 'cpu' # half precision only supported on CUDA |
|
|
|
self.trtFlag_det = True ###是否采用TRT模型加速 |
|
|
|
# 行人模型 |
|
|
|
class PedestrianModel: |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载行人模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
self.trtFlag_det = True, # 检测模型是否采用TRT |
|
|
|
if self.trtFlag_det: |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
det_gpuname = '_3090_fp16.engine' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
det_gpuname = '_2080Ti_fp16.engine' |
|
|
|
elif 'A10' in gpu.name: |
|
|
|
det_gpuname = '_A10_fp16.engine' |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
else: |
|
|
|
det_gpuname = '.pt' |
|
|
|
par = { |
|
|
|
'device': device, # 显卡号,如果用TRT模型，只支持0（单显卡） |
|
|
|
'labelnames': "../AIlib2/weights/pedestrian/labelnames.json", # 检测类别对照表 |
|
|
|
'Detweights': "../AIlib2/weights/pedestrian/yolov5%s" % det_gpuname, # 检测模型路径 |
|
|
|
'slopeIndex': [], # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。 |
|
|
|
'seg_nclass': 2, # 分割模型类别数目，默认2类 |
|
|
|
'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目 |
|
|
|
'segPar': None, # 分割模型预处理参数 |
|
|
|
'Segweights': None, # 分割模型权重位置 |
|
|
|
'postFile': '../AIlib2/weights/pedestrian/para.json' # 后处理参数文件 |
|
|
|
} |
|
|
|
self.device = select_device(par.get('device')) |
|
|
|
self.half = device.type != 'cpu' # half precision only supported on CUDA |
|
|
|
Detweights = par.get('Detweights') |
|
|
|
if self.trtFlag_det: |
|
|
|
detweights="../AIlib/weights/road/best_%s_fp16.engine" % gpuname |
|
|
|
logger = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(detweights, "rb") as f, trt.Runtime(logger) as runtime: |
|
|
|
self.model = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
print("1111111111111111111111111111111111111111111111111111") |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.model=runtime.deserialize_cuda_engine(f.read()) |
|
|
|
print('####load TRT model :%s'%(Detweights)) |
|
|
|
else: |
|
|
|
detweights="../AIlib/weights/road/best.pt" |
|
|
|
self.model = attempt_load(detweights, map_location=self.device) # load FP32 model |
|
|
|
self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model |
|
|
|
if self.half: |
|
|
|
self.model.half() |
|
|
|
print("2222222222222222222222222222222222222222222222222222") |
|
|
|
self.names = get_labelnames(par.get('labelnames')) |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile')) |
|
|
|
self.digitFont = None |
|
|
|
self.label_arraylist = None |
|
|
|
self.segmodel = None |
|
|
|
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres |
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres |
|
|
|
def process(self, frame, width=1920): |
|
|
|
if self.label_arraylist is None: |
|
|
|
fontsize = int(width / 1920 * 40) |
|
|
|
@@ -194,30 +321,229 @@ class RFModel(Model): |
|
|
|
waterLineWidth = 1 |
|
|
|
if width >= 960: |
|
|
|
line_thickness = int(round(width / 1920 * 3) - 1) |
|
|
|
boxLine_thickness = int(round(width / 1920 * 2)) |
|
|
|
boxLine_thickness = int(round(width / 1920 * 3)) |
|
|
|
waterLineWidth = int(round(width / 1920 * 3)) |
|
|
|
numFontSize = float(format(width / 1920 * 1.1, '.1f')) |
|
|
|
numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小 |
|
|
|
self.digitFont = {'line_thickness': line_thickness, |
|
|
|
'boxLine_thickness': boxLine_thickness, |
|
|
|
'fontSize': numFontSize, |
|
|
|
'waterLineColor': (0, 255, 255), |
|
|
|
'waterLineWidth': waterLineWidth} |
|
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize, |
|
|
|
fontpath="../AIlib2/conf/platech.ttf") |
|
|
|
return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist, |
|
|
|
self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres, |
|
|
|
self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det) |
|
|
|
|
|
|
|
|
|
|
|
# 烟火模型 |
|
|
|
class SmogfireModel: |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载烟火模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
self.trtFlag_det = True, # 检测模型是否采用TRT |
|
|
|
if self.trtFlag_det: |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
det_gpuname = '_3090_fp16.engine' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
det_gpuname = '_2080Ti_fp16.engine' |
|
|
|
elif 'A10' in gpu.name: |
|
|
|
det_gpuname = '_A10_fp16.engine' |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
else: |
|
|
|
det_gpuname = '.pt' |
|
|
|
par = { |
|
|
|
'device': device, # 显卡号,如果用TRT模型，只支持0（单显卡） |
|
|
|
'labelnames': "../AIlib2/weights/smogfire/labelnames.json", # 检测类别对照表 |
|
|
|
'Detweights': "../AIlib2/weights/smogfire/yolov5%s" % det_gpuname, # 检测模型路径 |
|
|
|
'slopeIndex': [], # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。 |
|
|
|
'seg_nclass': 2, # 分割模型类别数目，默认2类 |
|
|
|
'segRegionCnt': 0, # 分割模型结果需要保留的等值线数目 |
|
|
|
'segPar': None, # 分割模型预处理参数 |
|
|
|
'Segweights': None, # 分割模型权重位置 |
|
|
|
'postFile': '../AIlib2/weights/smogfire/para.json' # 后处理参数文件 |
|
|
|
} |
|
|
|
self.device = select_device(par.get('device')) |
|
|
|
self.half = device.type != 'cpu' # half precision only supported on CUDA |
|
|
|
Detweights = par.get('Detweights') |
|
|
|
if self.trtFlag_det: |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.model=runtime.deserialize_cuda_engine(f.read()) |
|
|
|
print('####load TRT model :%s'%(Detweights)) |
|
|
|
else: |
|
|
|
self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model |
|
|
|
if self.half: |
|
|
|
self.model.half() |
|
|
|
self.names = get_labelnames(par.get('labelnames')) |
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile')) |
|
|
|
self.digitFont = None |
|
|
|
self.label_arraylist = None |
|
|
|
self.segmodel = None |
|
|
|
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres |
|
|
|
def process(self, frame, width=1920): |
|
|
|
if self.label_arraylist is None: |
|
|
|
fontsize = int(width / 1920 * 40) |
|
|
|
line_thickness = 1 |
|
|
|
boxLine_thickness = 1 |
|
|
|
waterLineWidth = 1 |
|
|
|
if width >= 960: |
|
|
|
line_thickness = int(round(width / 1920 * 3) - 1) |
|
|
|
boxLine_thickness = int(round(width / 1920 * 3)) |
|
|
|
waterLineWidth = int(round(width / 1920 * 3)) |
|
|
|
numFontSize = float(format(width / 1920 * 1.1, '.1f')) # 文字大小 |
|
|
|
self.digitFont = {'line_thickness': line_thickness, |
|
|
|
'boxLine_thickness': boxLine_thickness, |
|
|
|
'fontSize': numFontSize, |
|
|
|
'waterLineColor': (0, 255, 255), |
|
|
|
'waterLineWidth': waterLineWidth} |
|
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize, |
|
|
|
fontpath="../AIlib/conf/platech.ttf") |
|
|
|
fontpath="../AIlib2/conf/platech.ttf") |
|
|
|
return AI_process_forest([frame], self.model, self.segmodel, self.names, self.label_arraylist, |
|
|
|
self.rainbows, self.half, self.device, self.conf_thres, self.iou_thres, |
|
|
|
self.allowedList, font=self.digitFont, trtFlag_det=self.trtFlag_det) |
|
|
|
|
|
|
|
# 交通模型 |
|
|
|
class RFModel: |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载交通模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
trtFlag_det = True, # 检测模型是否采用TRT |
|
|
|
trtFlag_seg = True, # 分割模型是否采用TRT |
|
|
|
if trtFlag_det and trtFlag_seg: |
|
|
|
gpu = get_all_gpu_ids()[int(device)] |
|
|
|
if '3090' in gpu.name: |
|
|
|
det_gpuname = '_3090_fp16.engine' |
|
|
|
seg_gpuname = '_3090_fp16.engine' |
|
|
|
elif '2080' in gpu.name: |
|
|
|
det_gpuname = '_2080Ti_fp16.engine' |
|
|
|
seg_gpuname = '_2080Ti_fp16.engine' |
|
|
|
elif 'A10' in gpu.name: |
|
|
|
det_gpuname = '_A10_fp16.engine' |
|
|
|
seg_gpuname = '_A10_fp16.engine' |
|
|
|
else: |
|
|
|
raise Exception("未匹配到该GPU名称的模型, GPU: " + gpu.name) |
|
|
|
else: |
|
|
|
det_gpuname = '.pt' |
|
|
|
seg_gpuname = '.pth' |
|
|
|
par = { |
|
|
|
'device': device, # 显卡号,如果用TRT模型，只支持0（单显卡） |
|
|
|
'labelnames': "../AIlib2/weights/road/labelnames.json", # 检测类别对照表 |
|
|
|
'trtFlag_det': trtFlag_det, # 检测模型是否采用TRT |
|
|
|
'trtFlag_seg': trtFlag_seg, # 分割模型是否采用TRT |
|
|
|
'Detweights': "../AIlib2/weights/road/yolov5%s" % det_gpuname, # 检测模型路径 |
|
|
|
'slopeIndex': [], # 岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。 |
|
|
|
'seg_nclass': 2, # 分割模型类别数目，默认2类 |
|
|
|
'segRegionCnt': 2, # 分割模型结果需要保留的等值线数目 |
|
|
|
'segPar': { |
|
|
|
'modelSize': (640,360), |
|
|
|
'mean': (0.485, 0.456, 0.406), |
|
|
|
'std' : (0.229, 0.224, 0.225), |
|
|
|
'numpy': False, |
|
|
|
'RGB_convert_first': True |
|
|
|
}, # 分割模型预处理参数 |
|
|
|
'Segweights': '../AIlib2/weights/road/stdc_360X640%s' % seg_gpuname, # 分割模型权重位置 |
|
|
|
'postFile': '../AIlib2/weights/road/para.json' # 后处理参数文件 |
|
|
|
} |
|
|
|
self.device = select_device(par.get('device')) |
|
|
|
self.names = get_labelnames(par.get('labelnames')) |
|
|
|
half = self.device.type != 'cpu' |
|
|
|
Detweights = par.get('Detweights') # 升级后的检测模型 |
|
|
|
if trtFlag_det: |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
print('############locad det model trt success#######') |
|
|
|
else: |
|
|
|
self.model = attempt_load(Detweights, map_location=self.device) # load FP32 model |
|
|
|
print('############locad det model pth success#######') |
|
|
|
if half: |
|
|
|
self.model.half() |
|
|
|
seg_nclass = par.get('seg_nclass') |
|
|
|
self.segPar = par.get('segPar') |
|
|
|
Segweights = par.get('Segweights') |
|
|
|
if trtFlag_seg: |
|
|
|
log = trt.Logger(trt.Logger.ERROR) |
|
|
|
with open(Segweights, "rb") as f, trt.Runtime(log) as runtime: |
|
|
|
self.segmodel = runtime.deserialize_cuda_engine(f.read()) # 输入trt本地文件，返回ICudaEngine对象 |
|
|
|
print('############locad seg model trt success#######') |
|
|
|
else: |
|
|
|
self.segmodel = SegModel(nclass=seg_nclass, weights=Segweights, device=self.device) |
|
|
|
print('############locad seg model pth success#######') |
|
|
|
|
|
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(par.get('postFile')) |
|
|
|
self.objectPar = { |
|
|
|
'half': half, |
|
|
|
'device': self.device, |
|
|
|
'conf_thres': self.conf_thres, |
|
|
|
'iou_thres': self.iou_thres, |
|
|
|
'allowedList': self.allowedList, |
|
|
|
'slopeIndex': par.get('slopeIndex'), |
|
|
|
'segRegionCnt': par.get('segRegionCnt'), |
|
|
|
'trtFlag_det': trtFlag_det, |
|
|
|
'trtFlag_seg': trtFlag_seg |
|
|
|
} |
|
|
|
self.label_arraylist = None |
|
|
|
self.digitFont = None |
|
|
|
|
|
|
|
def process(self, frame, width=1920): |
|
|
|
if self.label_arraylist is None: |
|
|
|
fontsize = int(width / 1920 * 40) |
|
|
|
line_thickness = 1 |
|
|
|
boxLine_thickness = 1 |
|
|
|
waterLineWidth = 1 |
|
|
|
if width >= 960: |
|
|
|
line_thickness = int(round(width / 1920 * 3) - 1) |
|
|
|
boxLine_thickness = int(round(width / 1920 * 2)) |
|
|
|
waterLineWidth = int(round(width / 1920 * 3)) |
|
|
|
numFontSize = float(format(width / 1920 * 1.1, '.1f')) |
|
|
|
|
|
|
|
self.digitFont = {'line_thickness': line_thickness, |
|
|
|
'boxLine_thickness': boxLine_thickness, |
|
|
|
'fontSize': numFontSize, |
|
|
|
'waterLineColor': (0, 255, 255), |
|
|
|
'waterLineWidth': waterLineWidth} |
|
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=fontsize, |
|
|
|
fontpath="../AIlib2/conf/platech.ttf") |
|
|
|
return AI_process([frame], self.model, self.segmodel, self.names, self.label_arraylist, |
|
|
|
self.rainbows, objectPar=self.objectPar, font=self.digitFont, segPar=self.segPar) |
|
|
|
|
|
|
|
|
|
|
|
# 车牌、健康码、行程码分割模型 |
|
|
|
# 健康码、行程码分割模型 |
|
|
|
class IMModel: |
|
|
|
def __init__(self, device, allowedList=None): |
|
|
|
self.order = allowedList[0] |
|
|
|
if int(allowedList[0]) == 3: |
|
|
|
self.img_type = 'plate' ## code,plate |
|
|
|
else: |
|
|
|
self.img_type = 'code' |
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载防疫模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
self.img_type = 'code' |
|
|
|
self.par = {'code': {'weights': '../AIlib2/weights/jkm/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10}, |
|
|
|
'plate': {'weights': '../AIlib2/weights/jkm/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1}, |
|
|
|
'conf_thres': 0.4, |
|
|
|
'iou_thres': 0.45, |
|
|
|
'device': 'cuda:%s' % device, |
|
|
|
'plate_dilate': (0.5, 0.3) |
|
|
|
} |
|
|
|
self.device = torch.device(self.par['device']) |
|
|
|
self.model = torch.jit.load(self.par[self.img_type]['weights']) |
|
|
|
|
|
|
|
def process(self, frame): |
|
|
|
img, padInfos = pre_process(frame, self.device) ##预处理 |
|
|
|
pred = self.model(img) ##模型推理 |
|
|
|
boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'], |
|
|
|
iou_thres=self.par['iou_thres'], nc=self.par[self.img_type]['nc']) # 后处理 |
|
|
|
dataBack = get_return_data(frame, boxes, modelType=self.img_type, plate_dilate=self.par['plate_dilate']) |
|
|
|
return dataBack |
|
|
|
|
|
|
|
|
|
|
|
# 车牌分割模型 |
|
|
|
class PlateMModel: |
|
|
|
|
|
|
|
def __init__(self, device, allowedList=None, logger=None, requestId=None): |
|
|
|
logger.info("########################加载车牌模型########################, requestId:{}", requestId) |
|
|
|
self.allowedList = allowedList |
|
|
|
self.img_type = 'plate' ## code,plate |
|
|
|
self.par = {'code': {'weights': '../AIlib/weights/jkm/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10}, |
|
|
|
'plate': {'weights': '../AIlib/weights/jkm/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1}, |
|
|
|
'conf_thres': 0.4, |
|
|
|
@@ -235,4 +561,4 @@ class IMModel: |
|
|
|
boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'], |
|
|
|
iou_thres=self.par['iou_thres'], nc=self.par[self.img_type]['nc']) # 后处理 |
|
|
|
dataBack = get_return_data(frame, boxes, modelType=self.img_type, plate_dilate=self.par['plate_dilate']) |
|
|
|
return dataBack |
|
|
|
return dataBack |
