1)新增M029 火焰面积 2）算法支持按类过滤 3）算法支持按置信度过滤 4）其他优化

AI.py

@@ -5,11 +5,13 @@ from segutils.trtUtils import segtrtEval,yolov5Trtforward,OcrTrtForward
 from segutils.trafficUtils import tracfficAccidentMixFunction
 
 from utils.torch_utils import select_device
-from utilsK.queRiver import get_labelnames,get_label_arrays,post_process_,img_pad,draw_painting_joint,detectDraw,getDetections,getDetectionsFromPreds
+from utilsK.queRiver import get_labelnames, img_pad, getDetections, getDetectionsFromPreds, scale_back
 from utilsK.jkmUtils import pre_process, post_process, get_return_data
 from trackUtils.sort import moving_average_wang
 
 from utils.datasets import letterbox
+from utils.general import non_max_suppression, scale_coords,xyxy2xywh,overlap_box_suppression
+from utils.plots import draw_painting_joint,get_label_arrays
 import numpy as np
 import torch
 import math
@@ -18,6 +20,7 @@ import torch.nn.functional as F
 from copy import deepcopy
 from scipy import interpolate
 import glob
+from loguru import logger
 
 def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=['.MP4','.mp4','.avi']):
     imgpaths=[];###获取文件里所有的图像
@@ -35,7 +38,7 @@ def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=[
     print('%s: test Images:%d , test videos:%d '%(impth, len(imgpaths), len(videopaths)))    
     return imgpaths,videopaths
 
-def xywh2xyxy(box,iW=None,iH=None):
+def xywh2xy(box,iW=None,iH=None):
     xc,yc,w,h = box[0:4]
     x0 =max(0, xc-w/2.0)
     x1 =min(1, xc+w/2.0)
@@ -73,13 +76,15 @@ def score_filter_byClass(pdetections,score_para_2nd):
             ret.append(det)
     return ret
 # 按类过滤
-def filter_byClass(pdetections,allowedList):
+def filter_byClass(pdetections, fiterList):
     ret = []
     for det in pdetections:
         score, cls = det[4], det[5]
-        if int(cls) in allowedList:
+        if int(cls) in fiterList:
-            ret.append(det)
+            continue
-        elif str(int(cls)) in allowedList:
+        elif str(int(cls)) in fiterList:
+            continue
+        else:
             ret.append(det)
 
     return ret
@@ -111,8 +116,89 @@ def plat_format(ocr):
 
     return label.upper()
 
-def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False,'score_byClass':{x:0.1 for x in range(30)} }, font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ,segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},mode='others',postPar=None):
+def post_process_det(pred,padInfos,img,im0s,conf_thres,iou_thres,label_arraylist,rainbows,font,score_byClass,fiterList,ovlap_thres=None):
+    time0 = time.time()
+    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)
+    if ovlap_thres:
+        pred = overlap_box_suppression(pred, ovlap_thres)
+    time1 = time.time()
+    det = pred[0]  ###一次检测一张图片
+    det_xywh = [];
+    im0 = im0s.copy()
+    #im0 = im0s[0]
+    if len(det) > 0:
+        # Rescale boxes from img_size to im0 size
+        if not padInfos:
+            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
+        else:
+            # print('####line131:',det[:, :])
+            det[:, :4] = scale_back(det[:, :4], padInfos).round()
 
+
+    for *xyxy, conf, cls in reversed(det):
+        cls_c = cls.cpu().numpy()
+        conf_c = conf.cpu().numpy()
+        tt = [int(x.cpu()) for x in xyxy]
+        if fiterList:
+            if int(cls) in fiterList:  ###如果不是所需要的目标，则不显示
+                continue
+        if score_byClass:
+            if  int(cls) in score_byClass.keys():
+                if conf < score_byClass[int(cls)]:
+                    continue
+        line = [*tt, float(conf_c), float(cls_c)]  # label format
+        det_xywh.append(line)
+    time2 = time.time()
+    strout='nms:%s ,detDraw:%s '%(get_ms(time0,time1), get_ms(time1,time2)   )
+    return [im0s[0],im0s[0], det_xywh, 10],strout
+
+def post_process_seg(im0s,segmodel,boxes,ksize):
+    time0 = time.time()
+    im0 = im0s[0].copy()
+    segmodel.set_image(im0s[0])
+    # # 创建一个空白掩码用于保存所有火焰
+    # combined_mask = np.zeros((im0.shape[0], im0.shape[1]), dtype=np.uint8)
+    # # 创建边缘可视化图像
+    # edge_image = np.zeros_like(im0s[0])
+    # 处理每个火焰检测框
+    det_xywhP = []
+    for box in boxes:
+        x_min, y_min, x_max, y_max = box[:4]
+        # 转换为SAM需要的格式
+        input_box = np.array([x_min, y_min, x_max, y_max])
+
+        # 使用框提示进行分割
+        masks, _, _ = segmodel.predict(
+            box=input_box,
+            multimask_output=False  # 只返回最佳掩码
+        )
+
+        # 获取分割掩码
+        flame_mask = masks[0].astype(np.uint8)
+        # 使用形态学操作填充小孔洞
+        filled_mask = cv2.morphologyEx(flame_mask, cv2.MORPH_CLOSE, ksize)
+        # 查找所有轮廓（包括内部小点）
+        contours, _ = cv2.findContours(filled_mask.astype(np.uint8),
+                                    cv2.RETR_TREE,
+                                    cv2.CHAIN_APPROX_SIMPLE)
+        if len(contours) == 0:
+            continue
+        largest_contour = max(contours, key=cv2.contourArea)
+        # 通过轮廓填充。
+        #cv2.drawContours(im0, [largest_contour], -1, (0, 0, 255), 2)
+        box.append(largest_contour)
+        det_xywhP.append(box)
+    time1 = time.time()
+    strout = 'segDraw:%s ' % get_ms(time0, time1)
+    return [im0s[0], im0s[0], det_xywhP, 10], strout
+
+def AI_process(im0s, model, segmodel, names, label_arraylist, rainbows,
+               objectPar={'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
+                          'segRegionCnt': 1, 'trtFlag_det': False,'trtFlag_seg': False,'score_byClass':None,'fiterList':[]},
+               font={'line_thickness': None, 'fontSize': None, 'boxLine_thickness': None,
+                     'waterLineColor': (0, 255, 255), 'waterLineWidth': 3},
+               segPar={'modelSize': (640, 360), 'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225),
+                       'numpy': False, 'RGB_convert_first': True}, mode='others', postPar=None):
     # 输入参数
     #   im0s---原始图像列表
     #   model---检测模型，segmodel---分割模型（如若没有用到，则为None）
@@ -126,14 +212,11 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
     #   #strout---统计AI处理个环节的时间
     # Letterbox
 
-    half,device,conf_thres,iou_thres,allowedList = objectPar['half'],objectPar['device'],objectPar['conf_thres'],objectPar['iou_thres'],objectPar['allowedList']
+    half, device, conf_thres, iou_thres, fiterList,score_byClass = objectPar['half'], objectPar['device'], objectPar['conf_thres'], \
+    objectPar['iou_thres'], objectPar['fiterList'], objectPar['score_byClass']
 
-    trtFlag_det,trtFlag_seg,segRegionCnt =   objectPar['trtFlag_det'],objectPar['trtFlag_seg'],objectPar['segRegionCnt']
+    trtFlag_det, trtFlag_seg, segRegionCnt = objectPar['trtFlag_det'], objectPar['trtFlag_seg'], objectPar[
-    if 'ovlap_thres_crossCategory' in objectPar.keys(): ovlap_thres = objectPar['ovlap_thres_crossCategory']
+        'segRegionCnt']
-    else: ovlap_thres = None
-
-    if 'score_byClass' in objectPar.keys(): score_byClass = objectPar['score_byClass']
-    else: score_byClass = None
 
     time0 = time.time()
     if trtFlag_det:
@@ -170,9 +253,8 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
     time2=time.time()
 
 
-    p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=ovlap_thres,padInfos=padInfos) 
+    p_result, timeOut = getDetectionsFromPreds(pred, img, im0s[0], conf_thres=conf_thres, iou_thres=iou_thres,
-    if score_byClass:
+                                               ovlap_thres=None, padInfos=padInfos)
-        p_result[2] = score_filter_byClass(p_result[2],score_byClass)
     # if  mode=='highWay3.0':
     # if segmodel:
     if segPar and segPar['mixFunction']['function']:
@@ -188,16 +270,23 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
     else:
         timeMixPost = ':0 ms'
     # print('#### line121: segstr:%s    timeMixPost:%s    timeOut:%s'%( segstr.strip(), timeMixPost,timeOut  ))
-    time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s  ,timeMixPost:%s '%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr.strip(),timeMixPost )
+    time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s  ,timeMixPost:%s ' % (
-    if allowedList:
+    (time01 - time0) * 1000, (time1 - time01) * 1000, (time2 - time1) * 1000, timeOut, segstr.strip(), timeMixPost)
-        p_result[2] = filter_byClass(p_result[2],allowedList)
+    if fiterList:
+        p_result[2] = filter_byClass(p_result[2], fiterList)
+
+    if score_byClass:
+        p_result[2] = score_filter_byClass(p_result[2], score_byClass)
 
     print('-' * 10, p_result[2])
     return p_result, time_info
+
+
 def default_mix(predlist, par):
     return predlist[0], ''
-def AI_process_N(im0s,modelList,postProcess):
 
+
+def AI_process_N(im0s, modelList, postProcess,score_byClass=None,fiterList=[]):
     # 输入参数
     ##   im0s---原始图像列表
     ##   modelList--所有的模型
@@ -221,7 +310,13 @@ def AI_process_N(im0s,modelList,postProcess):
 
     # ret就是混合处理后的结果
     ret = mixFunction(predsList, postProcess['pars'])
-    return ret[0],timeInfos+ret[1]
+    det = ret[0]
+    if fiterList:
+        det = filter_byClass(det, fiterList)
+    if score_byClass:
+        det = score_filter_byClass(det, score_byClass)
+
+    return det, timeInfos + ret[1]
 def getMaxScoreWords(detRets0):
     maxScore=-1;maxId=0
     for i,detRet in enumerate(detRets0):
@@ -230,7 +325,7 @@ def getMaxScoreWords(detRets0):
             maxScore = detRet[4]
     return maxId
 
-def AI_process_C(im0s,modelList,postProcess):
+def AI_process_C(im0s, modelList, postProcess,score_byClass,fiterList):
     # 函数定制的原因：
     ## 之前模型处理流是
     ## 图片---> 模型1-->result1；图片---> 模型2->result2；[result1,result2]--->后处理函数
@@ -293,6 +388,11 @@ def AI_process_C(im0s,modelList,postProcess):
             ocrInfo = detRets1[0][1]
         print('  _detRets0_obj:{} _detRets0_others:{} '.format(_detRets0_obj, _detRets0_others))
         rets = _detRets0_obj + _detRets0_others
+
+        if fiterList:
+            rets = filter_byClass(rets, fiterList)
+        if score_byClass:
+            rets = score_filter_byClass(rets, score_byClass)
         t3 = time.time()
         outInfos = 'total:%.1f ,where det:%.1f, ocr:%s' % ((t3 - t0) * 1000, (t1 - t0) * 1000, ocrInfo)
 
@@ -300,7 +400,11 @@ def AI_process_C(im0s,modelList,postProcess):
 
     return rets,outInfos
 
-def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=True,device=' cuda:0',conf_thres=0.25, iou_thres=0.45,allowedList=[0,1,2,3], font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ,trtFlag_det=False,SecNms=None):
+def AI_process_forest(im0s, model, segmodel, names, label_arraylist, rainbows, half=True, device=' cuda:0',
+                      conf_thres=0.25, iou_thres=0.45,
+                      font={'line_thickness': None, 'fontSize': None, 'boxLine_thickness': None,
+                            'waterLineColor': (0, 255, 255), 'waterLineWidth': 3}, trtFlag_det=False,
+                      SecNms=None,ksize=None,score_byClass=None,fiterList=[]):
     # 输入参数
     #   im0s---原始图像列表
     #   model---检测模型，segmodel---分割模型（如若没有用到，则为None）
@@ -329,24 +433,20 @@ def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=Tr
     img = img.half() if half else img.float()  # uint8 to fp16/32
 
     img /= 255.0  # 0 - 255 to 0.0 - 1.0
-    if segmodel:
+
-        seg_pred,segstr = segmodel.eval(im0s[0] )
-        segFlag=True
-    else:
-        seg_pred = None;segFlag=False
     time1 = time.time()
     pred = yolov5Trtforward(model, img) if trtFlag_det else model(img, augment=False)[0]
 
+    p_result, timeOut = post_process_det(pred,padInfos,img,im0s,conf_thres,iou_thres,label_arraylist,rainbows,font,score_byClass,fiterList)
+    if segmodel and len(p_result[2])>0:
+        segmodel.set_image(im0s[0])
+        p_result, timeOut = post_process_seg(im0s,segmodel,p_result[2],ksize)
 
     time2 = time.time()
-    datas = [[''], img, im0s, None,pred,seg_pred,10]
-
-    ObjectPar={ 'object_config':allowedList, 'slopeIndex':[] ,'segmodel':segFlag,'segRegionCnt':0 }
-    p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos,ovlap_thres=SecNms)
-    #print('###line274:',p_result[2])
-    #p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,object_config=allowedList,segmodel=segFlag,font=font,padInfos=padInfos)
     time_info = 'letterbox:%.1f, infer:%.1f, ' % ((time1 - time0) * 1000, (time2 - time1) * 1000)
     return p_result, time_info + timeOut
+
+
 def AI_det_track(im0s_in, modelPar, processPar, sort_tracker, segPar=None):
     im0s, iframe = im0s_in[0], im0s_in[1]
     model = modelPar['det_Model']
@@ -709,13 +809,18 @@ def AI_process_Ocr(im0s,modelList,device,detpar):
     timeMixPost = ':0 ms'
     new_device = torch.device(device)
     time0 = time.time()
+
     img, padInfos = pre_process(im0s[0], new_device)
     ocrModel = modelList[1]
     time1 = time.time()
+    if not detpar['trtFlag_det']:
         preds, timeOut = modelList[0].eval(img)
-    time2 = time.time()
         boxes = post_process(preds, padInfos, device, conf_thres=detpar['conf_thres'], iou_thres=detpar['iou_thres'],
                              nc=detpar['nc'])  # 后处理
+    else:
+        boxes, timeOut = modelList[0].eval(im0s[0])
+    time2 = time.time()
+
     imagePatches = [im0s[0][int(x[1]):int(x[3]), int(x[0]):int(x[2])] for x in boxes]
 
     detRets1 = [ocrModel.eval(patch) for patch in imagePatches]