AIlib2/utilsK/crowdGather.py

import sys
from pathlib import Path
import math
import cv2
import numpy as np
import torch
import math
import time

FILE = Path(__file__).absolute()
#sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path

def calculate_distance(point1, point2):
    """计算两个点之间的欧氏距离"""
    point= center_coordinate(point1)
    point=np.array(point)
    other_point = center_coordinate(point2)
    other_point = np.array(other_point)
    return np.linalg.norm(point - other_point)

def find_clusters(preds, min_distance):
    """按照最小距离将点分成簇"""
    points=preds
    points=np.array(points)
    clusters = []
    used_points = set()
    for i, point in enumerate(points):
        if i not in used_points:  # 如果该点未被使用过
            cluster = [point]
            used_points.add(i)
            # 寻找与该点距离小于等于min_distance的其他点
            for j, other_point in enumerate(points):
                if j not in used_points:
                    if all(calculate_distance(point, other_point) <= min_distance
                           for point in cluster):
                        cluster.append(other_point)
                        used_points.add(j)
            clusters.append(cluster)
    return clusters

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 get_bounding_rectangle(rectangles):
    '''
    通过输入多个矩形的对角坐标，得到这几个矩形的外包矩形对角坐标
    输入：点簇列表 （嵌套列表）
    输出：多个矩形的外包矩形对角坐标 （列表）
    '''
    min_x, max_x, min_y, max_y = float('inf'), float('-inf'), float('inf'), float('-inf')
    for rect in rectangles:
        x1, y1, x2, y2,c1,t1 = rect
        min_x = min(min_x, min(x1, x2))
        max_x = max(max_x, max(x1, x2))
        min_y = min(min_y, min(y1, y2))
        max_y = max(max_y, max(y1, y2))
    return [min_x, min_y, max_x, max_y]

def calculate_score(input_value):
    '''
    计算人群聚集置信度，检测出3-10人内，按照0.85-1的上升趋势取值；
    当检测超过10人，直接判断分数为1.
    '''

    if input_value == 3:
        output_value=0.85
    elif input_value == 4:
        output_value=0.9
    elif 5<= input_value <=10:
        output_value = 0.9+(input_value-4)*0.015
    else:
        output_value=1
    return output_value


def gather_post_process(predsList, pars):
    '''
    后处理程序，针对检测出的pedestrian，进行人员聚集的算法检测，按照类别'crowd_people'增加predsList
    ①原类别：
    ['ForestSpot', 'PestTree', 'pedestrian', 'fire', 'smog','cloud']=[0,1,2,3,4,5]
    ②处理后的类别汇总：
    ['ForestSpot', 'PestTree', 'pedestrian', 'fire', 'smog','cloud','crowd_people']=[0,1,2,3,4,5,6]

    输入：
    preds 一张图像的检测结果，为嵌套列表，tensor，包括x_y_x_y_conf_class
    imgwidth,imgheight 图像的原始宽度及长度
    输出：检测结果（将其中未悬挂国旗的显示）
    '''
    t0=time.time()
    predsList = predsList[0]
    predsList = [x for x in predsList if int(x[5]) !=5 ]##把类别“云朵”去除
    # 1、过滤掉类别2以外的目标，只保留行人
    preds = [ x for x in  predsList  if int(x[5]) ==pars['pedestrianId']  ]
    
   
    if len(preds)<  pars['crowdThreshold']:
        return predsList,'gaher postTime:No gathering'
    preds = np.array(preds)
    longs = np.mean(np.max(preds[:,2:4]-preds[:,0:2]))
    distanceThreshold =  pars['distancePersonScale']*longs
    # 2、查找点簇
    clusters = find_clusters(preds, distanceThreshold)

    clusters_crowd = []
    # 3、输出点簇信息，点簇中数量超过阈值，判断人员聚集
    for i, cluster in enumerate(clusters):
        if len(cluster) >= pars['crowdThreshold']:  # 超过一定人数，即为人员聚集
            #print(f"Cluster {i + 1}: {len(cluster)} points")
            clusters_crowd.append(cluster)
    #print(clusters_crowd)

    # 4、根据得到的人员聚集点簇，合并其他类别检测结果
    for i in range(len(clusters_crowd)):
        xyxy = get_bounding_rectangle(clusters_crowd[i])  # 人群聚集包围框
        score = calculate_score(len(clusters_crowd[i]))  # 人群聚集置信度
        xyxy.append(score)  # 人群聚集置信度
        xyxy.append(pars['gatherId'])  # 人群聚集类别
        predsList.append(xyxy)

    # 5、输出最终类别，共7类，用于绘图显示
    output_predslist = predsList
    #print('craoGaher line131:',output_predslist)
    t1=time.time()
    
    return output_predslist,'gaher postTime:%.1f ms'%(  (t1-t0)*1000  )

if __name__ == "__main__":
    t1 = time.time()
    # 对应vendor1_20240529_99.jpg检测结果
    preds=[[224.19933, 148.30751, 278.19156, 199.87828,   0.87625,   2.00000],
        [362.67139, 161.25760, 417.72357, 211.51706,   0.86919,   2.00000],
        [437.00131, 256.19083, 487.88870, 307.72897,   0.85786,   2.00000],
        [442.64606, 335.78168, 493.75720, 371.41418,   0.85245,   2.00000],
        [324.58362, 256.18488, 357.72626, 294.08929,   0.84512,   2.00000],
        [343.59781, 301.06506, 371.04105, 350.01086,   0.84207,   2.00000],
        [301.35858, 210.64088, 332.64862, 250.78883,   0.84063,   2.00000],
        [406.02994, 216.91214, 439.44455, 249.26077,   0.83698,   2.00000],
        [321.53494,  99.68467, 354.67477, 135.53226,   0.82515,   2.00000],
        [253.97131, 202.65234, 302.06055, 233.30634,   0.81498,   2.00000],
        [365.62521,  66.42108, 442.02292, 127.37558,   0.79556,   1.00000]]
    #preds=torch.tensor(preds)    #返回的预测结果
    imgwidth=1920
    imgheight=1680
    pars={'imgSize':(imgwidth,imgheight),'pedestrianId':2,'crowdThreshold':4,'gatherId':6,'distancePersonScale':2.0}
    '''
    pedestrianId 为行人识别的类别；
    crowdThreshold为设置的判断人员聚集的人数阈值，默认4人为聚集
    distanceThreshold为设置的判断人员聚集的距离阈值，为了测试默认300像素内为聚集(可自行设置)
    '''
    yyy=gather_post_process(preds,pars)  #送入后处理函数
    t2 = time.time()
    ttt = t2 - t1
    print('时间', ttt * 1000)