'''
这个版本增加了船舶过滤功能
'''
import time
import numpy as np
import cv2




def center_coordinate(boundbxs):
    '''
    输入：两个对角坐标xyxy
    输出：矩形框重点坐标xy
    '''
    boundbxs_x1=boundbxs[0]
    boundbxs_y1=boundbxs[1]
    boundbxs_x2=boundbxs[2]
    boundbxs_y2=boundbxs[3]
    center_x=0.5*(boundbxs_x1+boundbxs_x2)
    center_y=0.5*(boundbxs_y1+boundbxs_y2)
    return center_x,center_y

def fourcorner_coordinate(boundbxs):
    '''
    输入：两个对角坐标xyxy
    输出：矩形框四个角点坐标，以contours顺序。
    '''
    boundbxs_x1=boundbxs[0]
    boundbxs_y1=boundbxs[1]
    boundbxs_x2=boundbxs[2]
    boundbxs_y2=boundbxs[3]
    wid=boundbxs_x2-boundbxs_x1
    hei=boundbxs_y2-boundbxs_y1
    boundbxs_x3=boundbxs_x1+wid
    boundbxs_y3=boundbxs_y1
    boundbxs_x4=boundbxs_x1
    boundbxs_y4 = boundbxs_y1+hei
    contours_rec=[[boundbxs_x1,boundbxs_y1],[boundbxs_x3,boundbxs_y3],[boundbxs_x2,boundbxs_y2],[boundbxs_x4,boundbxs_y4]]
    return contours_rec



def expand_rectangle(rec,imgSize,ex_width,ex_height):
    '''
    矩形框外扩，且不超过图像范围
    输入：矩形框xyxy（左上和右下坐标），图像，外扩宽度大小，外扩高度大小
    输出：扩后的矩形框坐标xyxy
    '''
    #img_height=img.shape[0];img_width=img.shape[1]
    img_width,img_height = imgSize[0:2]
    #print('高、宽',img_height,img_width)
    x1=rec[0]
    y1=rec[1]
    x3=rec[2]
    y3=rec[3]

    x1=x1-ex_width if x1-ex_width >= 0 else 0
    y1=y1-ex_height if y1-ex_height >= 0 else 0
    x3=x3+ex_width if x3+ex_width <= img_width else img_width
    y3=y3+ex_height if y3+ex_height <=img_height else img_height
    xyxy=[x1,y1,x3,y3]

    return xyxy

def remove_simivalue(list1,list2):
    '''
    将list1中属于list2的元素都删除。
    输入：两个嵌套列表
    返回：嵌套列表
    '''
    list33=list1.copy()
    for i in range(len(list1)):
        for j in range(len(list2)):
            if list2[j] == list1[i]:
                # list33.pop(list1[i])
                list33.remove(list1[i])
    return list33

def remove_sameeleme_inalist(list3):
    '''
    将list3中重复嵌套列表元素删除。
    输入：嵌套列表
    返回：嵌套列表
    '''
    list3=list3
    list4=[]
    list4.append(list3[0])
    for dict in list3:
        k=0
        for item in list4:
            if dict!=item:
                k=k+1
            else:
                break
            if k==len(list4):
                list4.append(dict)
    return list4

def order_points(pts):
    ''' sort rectangle points by clockwise '''
    sort_x = pts[np.argsort(pts[:, 0]), :]

    Left = sort_x[:2, :]
    Right = sort_x[2:, :]
    # Left sort
    Left = Left[np.argsort(Left[:, 1])[::-1], :]
    # Right sort
    Right = Right[np.argsort(Right[:, 1]), :]
    return np.concatenate((Left, Right), axis=0)


def ms(t2,t1):
    
    return '%.1f' %( (t2-t1)*1000.0)
def illParking_postprocess(pred,cvMask,pars):
    #pred:直接预测结果，不要原图。预测结果[0,1,2,...],不是[车、T角点，L角点]
    #mask_cv:分割结果图，numpy格式(H,W),结果是int，[0,1,2,...]
    #pars: 其它参数,dict格式
    '''三个标签：车、T角点，L角点'''
    '''输入：落水人员的结果（类别+坐标）、原图
    
       过程：将车辆识别框外扩，并按contours形成区域。
            T角点与L角点的坐标合并为列表。
            判断每个车辆contours区域内有几个角点，少于2个则判断违停。
       返回：最终违停车辆标记结果图、违停车辆信息（坐标、类别、置信度）。
    '''
    #输入的是[cls,x0,y0,x1,y1,score]---> [x0,y0,x1,y1,cls,score]
    #输出的也是[cls,x0,y0,x1,y1,score]
    #pred = [  [ int(x[4])  ,*x[1:5], x[5] ] for x in pred]
    
    #pred = [[  *x[1:5],x[0], x[5] ] for x in pred]
    pred = [[  *x[0:4],x[5], x[4] ] for x in pred]

    ##统一格式
    imgSize=pars['imgSize']
    '''1、pred中车辆识别框形成列表，T角点与L角点形成列表'''
    tW1=time.time()
    init_vehicle=[]
    init_corner = []

    for i in range(len(pred)):
        #if pred[i][4]=='TCorner' or pred[i][4]=='LCorner':  #vehicle、TCorner、LCorner
        if pred[i][4]==1 or pred[i][4]==2:  #vehicle、TCorner、LCorner
            init_corner.append(pred[i])
        else:
            init_vehicle.append(pred[i])

    
    '''2、init_corner中心点坐标计算，并形成列表。'''
    tW2 = time.time()
    center_corner=[]
    for i in range(len(init_corner)):
        center_corner.append(center_coordinate(init_corner[i]))


    '''3、遍历每个车辆识别框，扩充矩形区域，将矩形区域形成contours，判断扩充区域内的。'''
    tW3 = time.time()
    final_weiting=[] #违停车辆列表
    '''遍历车辆列表，扩大矩形框形成contours'''
    for i in range(len(init_vehicle)):
        boundbxs1=[init_vehicle[i][0],init_vehicle[i][1],init_vehicle[i][2],init_vehicle[i][3]]
        width_boundingbox=init_vehicle[i][2]-init_vehicle[i][0]  #框宽度
        height_boundingbox=init_vehicle[i][2] - init_vehicle[i][0]  #框长度
        #当框长大于宽，则是水平方向车辆；否则认为是竖向车辆
        if width_boundingbox>=height_boundingbox:
            ex_width=0.4*(init_vehicle[i][2]-init_vehicle[i][0]) #矩形扩充宽度，取车宽0.4倍   #膨胀系数小一些。角点设成1个。
            ex_height=0.2*(init_vehicle[i][2]-init_vehicle[i][0]) #矩形扩充宽度，取车长0.2倍
            boundbxs1 = expand_rectangle(boundbxs1, imgSize, ex_width, ex_height)  # 扩充后矩形对角坐标
        else:
            ex_width=0.2*(init_vehicle[i][2]-init_vehicle[i][0]) #竖向，不需要改变变量名称，将系数对换下就行。（坐标点顺序还是1234不变）
            ex_height=0.4*(init_vehicle[i][2]-init_vehicle[i][0]) #
            boundbxs1 = expand_rectangle(boundbxs1, imgSize, ex_width, ex_height)  # 扩充后矩形对角坐标
        contour_temp=fourcorner_coordinate(boundbxs1) #得到扩充后矩形框的contour
        contour_temp_=np.array(contour_temp)#contour转为array
        contour_temp_=np.float32(contour_temp_)

        '''遍历角点识别框中心坐标是否在contours内，在则计1'''
        zzz=0
        for j in range(len(center_corner)):
            flag = cv2.pointPolygonTest(contour_temp_, (center_corner[j][0], center_corner[j][1]), False) #若为False，会找点是否在内，外，或轮廓上(相应返回+1, -1, 0)。
            if flag==+1:
                zzz+=1
        '''contours框内小于等于1个角点，认为不在停车位内'''
        # if zzz<=1:
        if zzz<1:
            final_weiting.append(init_vehicle[i])
    #print('t7-t6',t7-t6)
    #print('final_weiting',final_weiting)

    '''4、绘制保存检违停车辆图像'''
    
    tW4=time.time()
    '''
    colors = Colors()
    if final_weiting is not None:
        for i in range(len(final_weiting)):
            lbl='illegal park'
            xyxy=[final_weiting[i][0],final_weiting[i][1],final_weiting[i][2],final_weiting[i][3]]
            c = int(5)
            plot_one_box(xyxy, _img_cv, label=lbl, color=colors(c, True), line_thickness=3)
    final_img=_img_cv
    '''
    tW5=time.time()
    # cv2.imwrite('final_result.png', _img_cv)

    
    timeStr = ' step1:%s  step2:%s  step3:%s  save:%s'%( ms(tW2,tW1), ms(tW3,tW2),ms(tW4,tW3), ms(tW5,tW4)   )
    
    #final_weiting-----[x0,y0,x1,y1,cls,score]
    #输出的也是outRe----[cls,x0,y0,x1,y1,score]
    
    #outRes =  [  [ 3 ,*x[0:4], x[5] ] for x in final_weiting]###违停用3表示
    
    outRes =  [  [ *x[0:4], x[5],3 ] for x in final_weiting]###违停用3表示

    return outRes,timeStr #返回最终绘制的结果图、违停车辆（坐标、类别、置信度）

def illParking_postprocess_N(predList,pars):
    pred=predList[0]
    cvMask=None
    return illParking_postprocess(pred,cvMask,pars)
def AI_process(model,  args1,path1):
    '''对原图进行目标检测'''
    '''输入：检测模型、配置参数、路径
       返回：返回目标检测结果、原图像，
    '''
    '''检测图片'''
    t3=time.time()
    _img_cv = cv2.imread(path1)  # 将这里的送入yolov5
    t4 = time.time()
    pred = model.detect(_img_cv)  # 检测结果
    t5 = time.time()
    #print('t5-t4', t5-t4)
    #print('t4-t3', t4-t3)
    return pred, _img_cv   #返回目标检测结果、原图像

def main():

    '''配置参数'''

    args1={'cuda':'0','input_dir':'input_dir','output_dir':'output_dir'}

    dete_weights='weights/weiting20230727.pt'
    '''分割模型权重路径'''

    '''初始化目标检测模型'''
    model = Detector(dete_weights)

    names=['vehicle', 'TCorner', 'LCorner']
    t1=time.time()
    '''图像测试'''
    folders = os.listdir(args1['input_dir'])
    for i in range(len(folders)):
        path1 = args1['input_dir'] + '/' + folders[i]
        print('-'*100,path1)
        '''对原图进行目标检测'''
        pred, _img_cv=AI_process(model, args1,path1)
        H,W = _img_cv.shape[0:2]
        imgSize = (W,H);pars={'imgSize':imgSize}
        #preds = [[ names.index(x[4]),*x[0:4], float(x[5].cpu())  ] for x in pred[1]]
        preds = [[ *x[0:4], names.index(x[4]),float(x[5].cpu())  ] for x in pred[1]]
        # print('pred', pred)
        final_weiting,timeStr = illParking_postprocess(preds,None,pars)
        
        '''进入后处理，判断是否有违章停车'''
        #final_img,final_weiting=AI_postprocess(pred, _img_cv)
        #cv2.imwrite('./outdir/final_result'+str(i)+'.png', final_img)

    t2=time.time()
    print('耗时',t2-t1)

if __name__ == "__main__":
    main()

