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=[]
+    ret = []
     for det in pdetections:
-        score,cls = det[4],det[5]
+        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,9 +116,90 @@ 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,16 +212,13 @@ 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']
+    time0 = time.time()
-    else: score_byClass = None
-
-    time0=time.time() 
     if trtFlag_det:
         img, padInfos = img_pad(im0s[0], size=(640,640,3))  ;img = [img]
     else:
@@ -151,10 +234,10 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
     img = torch.from_numpy(img).to(device)
     img = img.half() if half else img.float()  # uint8 to fp16/32
     img /= 255.0
-    time01=time.time()
+    time01 = time.time()
 
     if segmodel:
-        seg_pred,segstr = segmodel.eval(im0s[0] )
+        seg_pred,segstr = segmodel.eval(im0s[0])
         segFlag=True
     else:
         seg_pred = None;segFlag=False;segstr='Not implemented'
@@ -170,11 +253,10 @@ 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  mode=='highWay3.0':
+    # if segmodel:
-    #if segmodel:
     if segPar and segPar['mixFunction']['function']:
 
         mixFunction = segPar['mixFunction']['function'];H,W = im0s[0].shape[0:2]
@@ -186,19 +268,26 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
         p_result.append(seg_pred)
 
     else:
-        timeMixPost=':0 ms'
+        timeMixPost = ':0 ms'
-    #print('#### line121: segstr:%s    timeMixPost:%s    timeOut:%s'%( segstr.strip(), timeMixPost,timeOut  ))
+    # 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)
 
-    print('-'*10,p_result[2])
+    if score_byClass:
-    return p_result,time_info
+        p_result[2] = score_filter_byClass(p_result[2], score_byClass)
-def default_mix(predlist,par):
-    return  predlist[0],''
-def AI_process_N(im0s,modelList,postProcess):
 
-    #输入参数
+    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,score_byClass=None,fiterList=[]):
+    # 输入参数
     ##   im0s---原始图像列表
     ##   modelList--所有的模型
     # postProcess--字典{},包括后处理函数，及其参数
@@ -206,22 +295,28 @@ def AI_process_N(im0s,modelList,postProcess):
     ##ret[0]--检测结果；
     ##ret[1]--时间信息
 
-    #modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果
+    # modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果
-    modelRets=[ model.eval(im0s[0]) for model in  modelList]
+    modelRets = [model.eval(im0s[0]) for model in modelList]
 
-    timeInfos = [ x[1] for x in modelRets]
+    timeInfos = [x[1] for x in modelRets]
-    timeInfos=''.join(timeInfos)
+    timeInfos = ''.join(timeInfos)
-    timeInfos=timeInfos
+    timeInfos = timeInfos
 
-    #postProcess['function']--后处理函数，输入的就是所有模型输出结果
+    # postProcess['function']--后处理函数，输入的就是所有模型输出结果
-    mixFunction =postProcess['function']
+    mixFunction = postProcess['function']
-    predsList = [ modelRet[0] for modelRet in modelRets ]
+    predsList = [modelRet[0] for modelRet in modelRets]
-    H,W = im0s[0].shape[0:2]
+    H, W = im0s[0].shape[0:2]
-    postProcess['pars']['imgSize'] = (W,H)
+    postProcess['pars']['imgSize'] = (W, H)
 
-    #ret就是混合处理后的结果
+    # ret就是混合处理后的结果
-    ret = mixFunction( predsList, postProcess['pars'])   
+    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,8 +325,8 @@ 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]--->后处理函数
     ## 本函数的处理流程是
@@ -287,21 +382,30 @@ def AI_process_C(im0s,modelList,postProcess):
             res_real = detRets1[0][0]
             res_real="".join(  list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),res_real)))
 
-            #detRets1[0][0]="".join(  list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),detRets1[0][0])))
+            # detRets1[0][0]="".join(  list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),detRets1[0][0])))
-            _detRets0_obj[maxId].append(res_real )
+            _detRets0_obj[maxId].append(res_real)
-            _detRets0_obj = [_detRets0_obj[maxId]]##只输出有OCR的那个船名结果
+            _detRets0_obj = [_detRets0_obj[maxId]]  ##只输出有OCR的那个船名结果
-            ocrInfo=detRets1[0][1]
+            ocrInfo = detRets1[0][1]
-        print(  '  _detRets0_obj:{} _detRets0_others:{} '.format( _detRets0_obj, _detRets0_others )  )
+        print('  _detRets0_obj:{} _detRets0_others:{} '.format(_detRets0_obj, _detRets0_others))
-        rets=_detRets0_obj+_detRets0_others
+        rets = _detRets0_obj + _detRets0_others
-        t3=time.time()
+
-        outInfos='total:%.1f ,where det:%.1f, ocr:%s'%(  (t3-t0)*1000, (t1-t0)*1000,  ocrInfo)
+        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)
 
         #print('###line233:',detRets1,detRets0 )
 
     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）
     #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout
@@ -313,7 +417,7 @@ def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=Tr
     #   #strout---统计AI处理个环节的时间
 
     # Letterbox
-    time0=time.time()
+    time0 = time.time()
     if trtFlag_det:
         img, padInfos = img_pad(im0s[0], size=(640,640,3))  ;img = [img]
     else:
@@ -329,26 +433,22 @@ 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] )
+    time1 = time.time()
-        segFlag=True
+    pred = yolov5Trtforward(model, img) if trtFlag_det else model(img, augment=False)[0]
-    else:
+
-        seg_pred = None;segFlag=False
+    p_result, timeOut = post_process_det(pred,padInfos,img,im0s,conf_thres,iou_thres,label_arraylist,rainbows,font,score_byClass,fiterList)
-    time1=time.time()
+    if segmodel and len(p_result[2])>0:
-    pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0]
+        segmodel.set_image(im0s[0])
+        p_result, timeOut = post_process_seg(im0s,segmodel,p_result[2],ksize)
+
+    time2 = time.time()
+    time_info = 'letterbox:%.1f, infer:%.1f, ' % ((time1 - time0) * 1000, (time2 - time1) * 1000)
+    return p_result, time_info + timeOut
 
 
-    time2=time.time()
+def AI_det_track(im0s_in, modelPar, processPar, sort_tracker, segPar=None):
-    datas = [[''], img, im0s, None,pred,seg_pred,10]
+    im0s, iframe = im0s_in[0], im0s_in[1]
-
-    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']
     segmodel = modelPar['seg_Model']
     half,device,conf_thres, iou_thres,trtFlag_det = processPar['half'], processPar['device'], processPar['conf_thres'], processPar['iou_thres'],processPar['trtFlag_det']
@@ -705,17 +805,22 @@ def ocr_process(pars):
     info_str= ('pre-process:%.2f  TRTforward:%.2f (%s)  postProcess:%2.f  decoder:%.2f, Total:%.2f , pred:%s'%(get_ms(time2,time1 ),get_ms(time3,time2 ),trtstr, get_ms(time4,time3 ), get_ms(time5,time4 ),  get_ms(time5,time1 ),  preds_str  )  )
     return  preds_str,info_str
 
-def AI_process_Ocr(im0s,modelList,device,detpar):
+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()
-    preds,timeOut = modelList[0].eval(img)
+    if not detpar['trtFlag_det']:
-    time2 = time.time()
+        preds, timeOut = modelList[0].eval(img)
         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]
@@ -730,7 +835,7 @@ def AI_process_Ocr(im0s,modelList,device,detpar):
     time_info = 'pre_process:%.1f, det:%.1f , ocr:%.1f  ,timeMixPost:%s ' % (
         (time1 - time0) * 1000, (time2 - time1) * 1000, (time3 - time2) * 1000, timeMixPost)
 
-    return [im0s[0],im0s[0],dets,0],time_info
+    return [im0s[0], im0s[0], dets, 0], time_info
 
 
 def AI_process_Crowd(im0s,model,device,postPar):