AIlib2/utilsK/spillUtils.py

import numpy as np
import time, cv2
from loguru import logger


def ms(t1, t0):
    return (t1 - t0) * 1000.0


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 mixSpillage_postprocess(preds, _mask_cv,pars=None):
    '''考虑船上人过滤'''
    '''输入：抛洒物的结果（类别+坐标）、原图、mask图像
       过程：获得mask的轮廓，判断抛洒物是否在轮廓内。
            在，则保留且绘制；不在，舍弃。
       返回：最终绘制的结果图、最终抛洒物（坐标、类别、置信度），
    '''
    '''1、最大分隔路面作为判断依据'''
    # zoom_factor=4 #缩小因子设置为4，考虑到numpy中分别遍历xy进行缩放耗时大。
    # speedroad = _mask_cv.copy()
    # speedroad = _mask_cv[speedroad==1]
    # _mask_cv[0] = _mask_cv[1]
    original_height = _mask_cv.shape[0]
    original_width = _mask_cv.shape[1]

    zoom_factor = original_width / 480.0
    
    zoom_height = int(original_height / zoom_factor)
    zoom_width = int(original_width / zoom_factor)

    _mask_cv = cv2.resize(_mask_cv, (zoom_width, zoom_height))  # 缩小原图,宽在前，高在后
    t4 = time.time()
    #print('+'*10,'_mask_cv shape信息',_mask_cv.shape)
    img_gray = cv2.cvtColor(_mask_cv, cv2.COLOR_BGR2GRAY) if len(_mask_cv.shape) == 3 else _mask_cv  #
    t5 = time.time()
    contours, thresh = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 寻找轮廓(多边界)
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, 2)
    contour_info = []
    for c in contours:
        contour_info.append((
            c,
            cv2.isContourConvex(c),
            cv2.contourArea(c),
        ))
    contour_info = sorted(contour_info, key=lambda c: c[2], reverse=True)
    t6 = time.time()

    # print('+'*10,'路面分隔信息',len(contour_info))
    '''新增模块：：如果路面为空，则返回原图、无抛洒物等。'''
    if contour_info == []:
        # final_img=_img_cv
        final_spillage_filterroad = []
        timeInfos = 'road is empty'

        return final_spillage_filterroad, timeInfos
    else:
        # print(contour_info[0])
        max_contour = contour_info[0]
        max_contour = max_contour[0] * zoom_factor  # contours恢复原图尺寸
        max_contour = max_contour.astype(np.int32)
        # print(max_contour)
    t7 = time.time()
    '''2.1、preds中spillage取出，car取出。'''
    init_spillage = []
    # init_car_per = []
    for i in range(len(preds)):
        if preds[i][5] == 0:
            init_spillage.append(preds[i])
        # else:
            # init_car_per.append(preds[i])
    #print('-'*10,'车辆',len(init_car_per))
    #print('+'*10,'抛洒物',len(init_spillage))
    #print('+'*10,'路面分隔信息',len(max_contour))
    # person

    # points = max_contour.reshape((-1, 1, 2))
    # cv2.polylines(image, [points], isClosed=True, color=(0, 255, 0), thickness=2)
    
    '''3、preds中spillage，通过1中路面过滤'''
    init_spillage_filterroad = init_spillage
    final_spillage_filterroad = []
    logger.info("车辆信息, max_contour: {}", max_contour)
    logger.info("车辆信息, init_spillage: {}", init_spillage)
    for i in range(len(init_spillage_filterroad)):
        center_x, center_y = center_coordinate(init_spillage_filterroad[i])
        # print('#'*20,'line176:',len(max_contour),np.array(max_contour).shape,(center_x, center_y))
        # 返回 1、-1 或 0，分别对应点在多边形内部、外部或边界上的情况
        flag = cv2.pointPolygonTest(max_contour, (int(center_x), int(center_y)),
                                    False)  # 若为False，会找点是否在内，外，或轮廓上(相应返回+1, -1, 0)。
        logger.info("车辆信息, flag: {}",flag)
        if flag == 1:
            final_spillage_filterroad.append(init_spillage_filterroad[i])
        else:
            pass
    t9 = time.time()

    
    timeInfos = ' findMaxroad:%.1f  releJudge:%.1f' % (ms(t6, t4), ms(t9, t6))

    return final_spillage_filterroad, timeInfos  # 返回最终绘制的结果图、最高速搞萨物（坐标、类别、置信度）