AIlib2/AI.py

import cv2,os,time,json
from models.experimental import attempt_load
from segutils.segmodel import SegModel,get_largest_contours
from segutils.trtUtils import segtrtEval,yolov5Trtforward,OcrTrtForward
from segutils.trafficUtils import trafficPostProcessing,colour_code_segmentation,get_label_info,trafficPostProcessingV2


from utils.torch_utils import select_device
from utilsK.queRiver import get_labelnames,get_label_arrays,post_process_,img_pad,draw_painting_joint

from utils.datasets import letterbox
import numpy as np
import torch
import math
from PIL import Image
import torch.nn.functional as F
from copy import deepcopy
def xywh2xyxy(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)
    y0=max(0, yc-h/2.0)
    y1=min(1,yc+h/2.0)
    if iW: x0,x1 = x0*iW,x1*iW
    if iH: y0,y1 = y0*iH,y1*iH
    return [x0,y0,x1,y1]    
    

def get_ms(t2,t1):
    return (t2-t1)*1000.0
def get_postProcess_para(parfile):
    with open(parfile) as fp:
        par = json.load(fp)
    assert 'post_process' in par.keys(), ' parfile has not key word:post_process'   
    parPost=par['post_process']  
    
    return parPost["conf_thres"],parPost["iou_thres"],parPost["classes"],parPost["rainbows"]  
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],'slopeIndex':[5,6,7],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False }, 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）
    #
    #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout
    #   [im0s[0],im0,det_xywh,iframe]中，
    #   im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。
    #   det_xywh--检测结果，是一个列表。
    #       其中每一个元素表示一个目标构成如：[float(cls_c), xc,yc,w,h, float(conf_c)]
    #       #cls_c--类别，如0,1,2,3;  xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间
    #   #strout---统计AI处理个环节的时间
    # Letterbox

    half,device,conf_thres,iou_thres,allowedList = objectPar['half'],objectPar['device'],objectPar['conf_thres'],objectPar['iou_thres'],objectPar['allowedList']
    slopeIndex, trtFlag_det,trtFlag_seg,segRegionCnt =    objectPar['slopeIndex'],objectPar['trtFlag_det'],objectPar['trtFlag_seg'],objectPar['segRegionCnt']
    
    time0=time.time() 
    if trtFlag_det:
        img, padInfos = img_pad(im0s[0], size=(640,640,3))  ;img = [img]
    else:
        img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
    # Stack
    img = np.stack(img, 0)
    # Convert
    img = img[:, :, :, ::-1].transpose(0, 3, 1, 2)  # BGR to RGB, to bsx3x416x416
    img = np.ascontiguousarray(img)  
    
    
    img = torch.from_numpy(img).to(device)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    time01=time.time()
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if segmodel:
        if trtFlag_seg:
           seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar)
        else:
            seg_pred,segstr = segmodel.eval(im0s[0] )
        segFlag=True
    else:
        seg_pred = None;segFlag=False;segstr='Not implemented'
    if mode=='highWay3.0':
        seg_pred_mulcls = seg_pred.copy()
        seg_pred = (seg_pred==1).astype(np.uint8) ###把路提取出来,路的类别是1
 
    time1=time.time()  
    if   trtFlag_det:  
        pred = yolov5Trtforward(model,img)
    else:
        pred = model(img,augment=False)[0]
       
    time2=time.time()
    datas = [[''], img, im0s, None,pred,seg_pred,10]        
    ObjectPar={ 'object_config':allowedList, 'slopeIndex':slopeIndex ,'segmodel':segFlag,'segRegionCnt':segRegionCnt }

    p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos)
   
    if  mode=='highWay3.0':
        assert postPar , ' postPar not implemented'
        label_info = get_label_info(postPar['label_csv'])  
        seg_pred=cv2.resize(seg_pred_mulcls,(   segPar['modelSize'][0]  ,  segPar['modelSize'] [1]) )
        imH,imW = im0s[0].shape[0:2]
        mmH,mmW  = seg_pred.shape[0:2]
        fx=mmW/imW;fy=mmH/imH
        det_coords=[]
        
        det_coords_original=[]
        for box in p_result[2]:  
            b_0 = box[1:5]
            b_0.insert(0,box[0]);b_0.append(box[5] )
            det_coords_original.append( b_0 )  
            if int(box[0]) != 1: continue  
            b=box[1:5]
            b[0],b[2] = b[0]*fx,b[2]*fx
            b[1],b[3] = b[1]*fy,b[3]*fy
            b=[int(x) for x in b]
            det_coords.append( (b[0],b[1])); det_coords.append(  (b[2],b[3])  )

        postPar['ZoomFactor']={'x':mmW/imW  ,'y':mmH/imH}
        postPar['mask']=seg_pred;postPar['det']=deepcopy(det_coords_original)
        postPar['label_info']=label_info
        tlist = list(postPar.keys()); tlist.sort()
        if len(det_coords_original)> 0:
           
            list8, image,time_infos = trafficPostProcessingV2(postPar) 
            
            
            Accident_results = np.array(list8,dtype=object)  
            acc_det=[]
            for bpoints in list8:
                if bpoints[9]>conf_thres:
                    xyxy=bpoints[1:5];xyxy=[int(x) for x in xyxy]
                    cls=9;conf=bpoints[9];
                    box_acc = [cls,*xyxy,conf]
                    acc_det.append(box_acc)
                    if cls  in allowedList:
                        p_result[1] = draw_painting_joint(xyxy,p_result[1],label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],font=font,socre_location="leftBottom")
                    
            det_coords_original.extend(acc_det)   
            p_result[2]=  deepcopy(det_coords_original)
        
        #print( '  time:',time_infos,'  results.shape:',Accident_results.shape, p_result[2])
    time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s'%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr )
    if mode=='highWay3.0':
        p_result.append(seg_pred_mulcls)

    return p_result,time_info
    
def AI_Seg_process(im0s,segmodel,digitWordFont,trtFlag_seg=True,segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},postPar= {'label_csv': './AIlib2/weights/conf/trafficAccident/class_dict.csv', 'speedRoadArea': 5100, 'vehicleArea': 100, 'speedRoadVehicleAngleMin': 15, 'speedRoadVehicleAngleMax': 75, 'vehicleLengthWidthThreshold': 4, 'vehicleSafeDistance': 7}):
    '''
    输入参数
    im0s---原始图像列表
    segmodel---分割模型，segmodel---分割模型（如若没有用到，则为None）
    digitWordFont--显示字体，数字等参数
    trtFlag_seg--模型是否是TRT格式
    segPar--分割模型的参数
    postPar--后处理参数
    输出
    seg_pred--返回语义分割的结果图（0,1,2...表示）
    img_draw--原图上带有矩形框的图
    segstr-----文本数据包括时间信息
    list1-----返回目标的坐标结果，每一个目标用[ cls, x0,y0,x1,y1,conf ]
    '''
    time1=time.time()
    H,W=im0s[0].shape[0:2]
    img_draw=im0s[0].copy()
    if trtFlag_seg:
        seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar)
    else:
        seg_pred,segstr = segmodel.eval(im0s[0] )
    time2 = time.time()
    label_info = get_label_info(postPar['label_csv'])    
    postPar['CCS']=colour_code_segmentation(seg_pred.copy(), label_info)
    postPar['sourceImageSize'] = im0s[0].shape[0:2]
    postPar['seg_pred_size'] = seg_pred.shape[0:2]

    list1,post_time_infos = trafficPostProcessing(postPar)   
    list2=[]
    cls=0
    label_arraylist=digitWordFont['label_arraylist']
    rainbows=digitWordFont['rainbows']
    for bpoints in list1:
        #print('###line104:',bpoints)
        bpoints=np.array(bpoints)
        x0=np.min( bpoints[:,0] )
        y0=np.min( bpoints[:,1] )
        x1=np.max( bpoints[:,0] )
        y1=np.max( bpoints[:,1] )
        conf= ((x0+x1)/W + (y0+y1)/H)/4.0;
        conf=1.0 - math.fabs((conf-0.5)/0.5)
        xyxy=[x0,y0,x1,y1]
        xyxy=[int(x+0.5)  for x in xyxy]
        #float(cls_c), *xywh, float(conf_c)]
        list2.append( [ cls, x0,y0,x1,y1,conf ] )
        img_draw = draw_painting_joint(xyxy,img_draw,label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],font=digitWordFont)
     
    
    segstr = 'segInfer:%.2f  %s '%( (time2-time1)*1000.0,post_time_infos    )
    
    return seg_pred,img_draw,segstr,list2
def AI_process_v2(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} ):
    #输入参数
    #   im0s---原始图像列表
    #   model---检测模型，segmodel---分割模型（如若没有用到，则为None）
    #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout
    #   [im0s[0],im0,det_xywh,iframe]中，
    #   im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。
    #   det_xywh--检测结果，是一个列表。
    #       其中每一个元素表示一个目标构成如：[float(cls_c), xc,yc,w,h, float(conf_c)]
    #       #cls_c--类别，如0,1,2,3;  xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间
    #   #strout---统计AI处理个环节的时间
    
    
    
     # Letterbox
    time0=time.time() 
    #img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s]
    
    img, padInfos = img_pad(im0s[0], size=(640,640,3))  ;img = [img]
    
    # Stack
    img = np.stack(img, 0)
    # Convert
    img = img[:, :, :, ::-1].transpose(0, 3, 1, 2)  # BGR to RGB, to bsx3x416x416
    img = np.ascontiguousarray(img)  
    
    
    img = torch.from_numpy(img).to(device)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    time01=time.time()
    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 = model(img,augment=False)
    time2=time.time()
    datas = [[''], img, im0s, None,pred,seg_pred,10]        

    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, seg:%.1f , infer:%.1f,%s, seginfo:%s'%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr )
    return p_result,time_info
    
        
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):
    #输入参数
    #   im0s---原始图像列表
    #   model---检测模型，segmodel---分割模型（如若没有用到，则为None）
    #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout
    #   [im0s[0],im0,det_xywh,iframe]中，
    #   im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。
    #   det_xywh--检测结果，是一个列表。
    #       其中每一个元素表示一个目标构成如：[float(cls_c), xc,yc,w,h, float(conf_c)]
    #       #cls_c--类别，如0,1,2,3;  xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间
    #   #strout---统计AI处理个环节的时间
    
     # Letterbox
    time0=time.time() 
    if trtFlag_det:
        img, padInfos = img_pad(im0s[0], size=(640,640,3))  ;img = [img]
    else:
        img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
    #img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s]
    # Stack
    img = np.stack(img, 0)
    # Convert
    img = img[:, :, :, ::-1].transpose(0, 3, 1, 2)  # BGR to RGB, to bsx3x416x416
    img = np.ascontiguousarray(img)  

    img = torch.from_numpy(img).to(device)
    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]
    
        
    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)
    
    #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 ocr_process(pars):

    img_patch,engine,context,converter,AlignCollate_normal,device=pars[0:6]
    time1 = time.time()
    img_tensor = AlignCollate_normal([ Image.fromarray(img_patch,'L') ])
    img_input = img_tensor.to('cuda:0')
    time2 = time.time()
    
    preds,trtstr=OcrTrtForward(engine,[img_input],context)
    time3 = time.time()
    
    batch_size = preds.size(0)
    preds_size = torch.IntTensor([preds.size(1)] * batch_size)

    ######## filter ignore_char, rebalance
    preds_prob = F.softmax(preds, dim=2)
    preds_prob = preds_prob.cpu().detach().numpy()
    pred_norm = preds_prob.sum(axis=2)
    preds_prob = preds_prob/np.expand_dims(pred_norm, axis=-1)
    preds_prob = torch.from_numpy(preds_prob).float().to(device)
    _, preds_index = preds_prob.max(2)
    preds_index = preds_index.view(-1)  
    time4 = time.time()
    preds_str = converter.decode_greedy(preds_index.data.cpu().detach().numpy(), preds_size.data)
    time5 = time.time()
    
    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 main():
    ##预先设置的参数
    device_='1'  ##选定模型，可选 cpu,'0','1'
    
    ##以下参数目前不可改
    Detweights = "weights/yolov5/class5/best_5classes.pt"
    seg_nclass = 2 
    Segweights = "weights/BiSeNet/checkpoint.pth"
    conf_thres,iou_thres,classes= 0.25,0.45,5
    labelnames = "weights/yolov5/class5/labelnames.json"
    rainbows = [ [0,0,255],[0,255,0],[255,0,0],[255,0,255],[255,255,0],[255,129,0],[255,0,127],[127,255,0],[0,255,127],[0,127,255],[127,0,255],[255,127,255],[255,255,127],[127,255,255],[0,255,255],[255,127,255],[127,255,255],  [0,127,0],[0,0,127],[0,255,255]]
    allowedList=[0,1,2,3]
    
    
    ##加载模型，准备好显示字符
    device = select_device(device_)
    names=get_labelnames(labelnames)
    label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="conf/platech.ttf")
    half = device.type != 'cpu'  # half precision only supported on CUDA
    model = attempt_load(Detweights, map_location=device)  # load FP32 model
    if half: model.half()
    segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
    
    
    
    ##图像测试
    #url='images/examples/20220624_响水河_12300_1621.jpg'
    impth = 'images/examples/'
    outpth = 'images/results/'
    folders = os.listdir(impth)
    for i in range(len(folders)):
        imgpath = os.path.join(impth, folders[i])
        im0s=[cv2.imread(imgpath)]
        time00 = time.time()
        p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,fontSize=1.0)
        time11 = time.time()
        image_array = p_result[1]
        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,image_array )
        print('----process:%s'%(folders[i]), (time.time() - time11) * 1000)

    
    
    
        
if __name__=="__main__":
    main()