diff --git a/AI.py b/AI.py index f8d860c..6ed26ba 100644 --- a/AI.py +++ b/AI.py @@ -735,30 +735,87 @@ def AI_process_Ocr(im0s,modelList,device,detpar): def AI_process_Crowd(im0s,model,device,postPar): timeMixPost = ':0 ms' - new_device = torch.device(device) time0 = time.time() - preds = model.eval(im0s[0]) + h, w = im0s[0].shape[:2] + Wrate = w // 128 * 128 / w + Hrate = h // 128 * 128 / h + detM = np.array([]) + if len(model) == 1: + detP = model[0].eval(im0s[0]) + p_postPar = postPar[0] + + else: + detP = model[1].eval(im0s[0]) + p_postPar = postPar[1] + detM = np.array(model[0].eval(im0s[0])[0]) + # d_postPar = postPar[0] time1 = time.time() - outputs_scores = torch.nn.functional.softmax(preds['pred_logits'], -1)[:, :, 1][0] + outputs_scores = torch.nn.functional.softmax(detP['pred_logits'], -1)[:, :, 1][0] - outputs_points = preds['pred_points'][0] + outputs_points = detP['pred_points'][0] + + detP = outputs_points[outputs_scores > p_postPar['conf']].detach().cpu().numpy() + detP[:, 0] = detP[:, 0] / Wrate + detP[:, 1] = detP[:, 1] / Hrate + + dets = [] + if not detM.size==0: + P = detP.shape[0] # 第二类边界框数量 + if P == 0: + time_info = 'No Valid object find' + return [im0s[0], im0s[0], dets, 0], time_info + x1 = detP[:, 0] - 5 + y1 = detP[:, 1] - 5 + x2 = detP[:, 0] + 5 + y2 = detP[:, 1] + 5 + detP = np.column_stack((x1, y1, x2, y2)) + detM = detM[:,:4] + a_x1, a_y1, a_x2, a_y2 = detM[:, 0], detM[:, 1], detM[:, 2], detM[:, 3] + b_x1, b_y1, b_x2, b_y2 = detP[:, 0], detP[:, 1], detP[:, 2], detP[:, 3] + + # 扩展维度以便广播 [N,1] + a_x1 = a_x1[:, np.newaxis] + a_y1 = a_y1[:, np.newaxis] + a_x2 = a_x2[:, np.newaxis] + a_y2 = a_y2[:, np.newaxis] + + # 计算相交区域坐标 [N,M] + inter_x1 = np.maximum(a_x1, b_x1) + inter_y1 = np.maximum(a_y1, b_y1) + inter_x2 = np.minimum(a_x2, b_x2) + inter_y2 = np.minimum(a_y2, b_y2) + # 计算相交区域面积 [N,M] + inter_area = np.maximum(0, inter_x2 - inter_x1) * np.maximum(0, inter_y2 - inter_y1) + + # 计算每个A框与多少个B框相交(相交面积>0) + counts = np.sum(inter_area > 0, axis=1).astype(int) + + # 构建结果数组:前4列是A框坐标,第5列是相交计数 + detM = np.hstack((detM, counts.reshape(-1, 1))) + detM = detM[detM[:, -1] != 0] + # 找出所有与A框相交的B框索引 + p_preds_index = np.any(inter_area > 0, axis=0) + p_indices = np.where(p_preds_index)[0] + # 提取与第一类相交的第二类边界框 + if len(p_indices) > 0: + detP = detP[p_indices] + detP = detP[:, :2] + detP[:, 0] = detP[:, 0] + 5 + detP[:, 1] = detP[:, 1] + 5 + + dets= [[detM.tolist(), detP.tolist()]] + else: + dets = [[detP.tolist()]] + + # for b in detM: + # b_label = '该建筑物下行人及数量:%d'%(int(b[4])) - points = outputs_points[outputs_scores > postPar['conf']].detach().cpu().numpy().tolist() - predict_cnt = int((outputs_scores > postPar['conf']).sum()) - #img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR) time2 = time.time() - # for p in points: - # img_to_draw = cv2.circle(img_to_draw, (int(p[0]), int(p[1])), line, (0, 0, 255), -1) - Calc_label = '当前人数: %d' % (predict_cnt) - - dets = [[Calc_label, points]] time_info = 'det:%.1f , post:%.1f ,timeMixPost:%s ' % ( - (time1 - time0) * 1000, (time2 - time1) * 1000, timeMixPost) + (time1 - time0) * 1000, (time2 - time1) * 1000, timeMixPost) - - - return [im0s[0],im0s[0],dets,0],time_info + return [im0s[0], im0s[0], dets, 0], time_info def main():