Spaces:
Running on Zero
Running on Zero
File size: 2,923 Bytes
09462dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import numpy as np
import copy
def get_bbox_area(bbox):
x1, y1, x2, y2 = bbox
return (x2 - x1) * (y2 - y1)
def check_single_human_requirements(det_result):
# filter results
if len(det_result) > 3 or len(det_result) == 0:
return False
elif len(det_result) == 1:
return True
elif len(det_result) > 1: # [2, 3]
bbox_areas = [get_bbox_area(bbox) for bbox in det_result]
# 获取最大 bbox 面积的索引
max_ind = max(range(len(bbox_areas)), key=lambda i: bbox_areas[i])
# 获取次大面积（需要排除 max_ind）
other_indices = [i for i in range(len(bbox_areas)) if i != max_ind]
second_max_area = max([bbox_areas[i] for i in other_indices])
max_area = bbox_areas[max_ind]
if max_area < 2 * second_max_area:
return False
else:
return True
def human_select(poses, det_results, multi_person):
new_poses = []
new_det_results = []
for pose, det_result in zip(poses, det_results):
if multi_person:
new_pose, new_det_result = get_multi_human(pose, det_result)
else:
new_pose, new_det_result = get_single_human(pose, det_result)
new_poses.append(new_pose)
new_det_results.append(new_det_result)
return new_poses, new_det_results
def get_single_human(pose, det_result):
if len(det_result) <= 1:
return pose, det_result
else:
bbox_areas = [get_bbox_area(bbox) for bbox in det_result]
max_ind = max(range(len(bbox_areas)), key=lambda i: bbox_areas[i])
pose_copy = copy.deepcopy(pose)
pose_copy['bodies']['candidate'] = pose_copy['bodies']['candidate'][max_ind:max_ind+1]
pose_copy['bodies']['subset'] = pose_copy['bodies']['subset'][max_ind:max_ind+1]
pose_copy['hands'] = pose_copy['hands'][2*max_ind:2*max_ind+2]
pose_copy['faces'] = pose_copy['faces'][max_ind:max_ind+1]
return pose_copy, det_result[max_ind:max_ind+1]
def check_multi_human_requirements(det_result):
# filter results
if len(det_result) < 2 or len(det_result) > 4: # 2-4个人
return False
else: # [3, 6]
bbox_areas = [get_bbox_area(bbox) for bbox in det_result]
# 获取最大 bbox 面积的索引
max_ind = max(range(len(bbox_areas)), key=lambda i: bbox_areas[i])
max_area = bbox_areas[max_ind]
# 选择面积大于等于最大面积50%的bbox
selected_indices = [i for i in range(len(bbox_areas)) if bbox_areas[i] >= 0.5 * max_area] # 包含max_ind
# 检查选中的bbox数量是否大于等于2
if len(selected_indices) >= 2:
return True
else:
return False
def get_multi_human(pose, det_result):
# 后续再筛比较好，后续从65帧里面筛的时候可以把背景里的人的筛掉
return pose, det_result |