Spaces:
Runtime error
Runtime error
File size: 6,613 Bytes
1cdc47e |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import json
import numpy as np
from tools.mpii_coco_h36m import coco_h36m
import os
h36m_coco_order = [9, 11, 14, 12, 15, 13, 16, 4, 1, 5, 2, 6, 3]
coco_order = [0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
num_person = 2
num_joints = 17
img_3d = 100.
ratio_2d_3d = 500.
def load_json(file_path):
with open(file_path, 'r') as fr:
video_info = json.load(fr)
label = video_info['label']
label_index = video_info['label_index']
num_frames = video_info['data'][-1]['frame_index']
keypoints = np.zeros((num_person, num_frames, num_joints, 2), dtype=np.float32)
scores = np.zeros((num_person, num_frames, num_joints), dtype=np.float32)
for frame_info in video_info['data']:
frame_index = frame_info['frame_index']
for index, skeleton_info in enumerate(frame_info['skeleton']):
pose = skeleton_info['pose']
score = skeleton_info['score']
bbox = skeleton_info['bbox']
if len(bbox) == 0 or index+1 > num_person:
continue
pose = np.asarray(pose, dtype=np.float32)
score = np.asarray(score, dtype=np.float32)
score = score.reshape(-1)
keypoints[index, frame_index-1] = pose
scores[index, frame_index-1] = score
return keypoints, scores, label, label_index
def h36m_coco_format(keypoints, scores):
assert len(keypoints.shape) == 4 and len(scores.shape) == 3
h36m_kpts = []
h36m_scores = []
valid_frames = []
for i in range(keypoints.shape[0]):
kpts = keypoints[i]
score = scores[i]
new_score = np.zeros_like(score, dtype=np.float32)
if np.sum(kpts) != 0.:
kpts, valid_frame = coco_h36m(kpts)
h36m_kpts.append(kpts)
valid_frames.append(valid_frame)
new_score[:, h36m_coco_order] = score[:, coco_order]
new_score[:, 0] = np.mean(score[:, [11, 12]], axis=1, dtype=np.float32)
new_score[:, 8] = np.mean(score[:, [5, 6]], axis=1, dtype=np.float32)
new_score[:, 7] = np.mean(new_score[:, [0, 8]], axis=1, dtype=np.float32)
new_score[:, 10] = np.mean(score[:, [1, 2, 3, 4]], axis=1, dtype=np.float32)
h36m_scores.append(new_score)
h36m_kpts = np.asarray(h36m_kpts, dtype=np.float32)
h36m_scores = np.asarray(h36m_scores, dtype=np.float32)
return h36m_kpts, h36m_scores, valid_frames
def revise_kpts(h36m_kpts, h36m_scores, valid_frames):
new_h36m_kpts = np.zeros_like(h36m_kpts)
for index, frames in enumerate(valid_frames):
kpts = h36m_kpts[index, frames]
score = h36m_scores[index, frames]
# threshold_score = score > 0.3
# if threshold_score.all():
# continue
index_frame = np.where(np.sum(score < 0.3, axis=1) > 0)[0]
for frame in index_frame:
less_threshold_joints = np.where(score[frame] < 0.3)[0]
intersect = [i for i in [2, 3, 5, 6] if i in less_threshold_joints]
if [2, 3, 5, 6] == intersect:
kpts[frame, [2, 3, 5, 6]] = kpts[frame, [1, 1, 4, 4]]
elif [2, 3, 6] == intersect:
kpts[frame, [2, 3, 6]] = kpts[frame, [1, 1, 5]]
elif [3, 5, 6] == intersect:
kpts[frame, [3, 5, 6]] = kpts[frame, [2, 4, 4]]
elif [3, 6] == intersect:
kpts[frame, [3, 6]] = kpts[frame, [2, 5]]
elif [3] == intersect:
kpts[frame, 3] = kpts[frame, 2]
elif [6] == intersect:
kpts[frame, 6] = kpts[frame, 5]
else:
continue
new_h36m_kpts[index, frames] = kpts
return new_h36m_kpts
def load_kpts_json(kpts_json):
keypoints, scores, label, label_index = load_json(kpts_json)
h36m_kpts, h36m_scores, valid_frames = h36m_coco_format(keypoints, scores)
re_kpts = revise_kpts(h36m_kpts, h36m_scores, valid_frames)
return re_kpts, valid_frames, scores, label, label_index
def revise_skes(prediction, re_kpts, valid_frames):
new_prediction = np.zeros((*re_kpts.shape[:-1], 3), dtype=np.float32)
for i, frames in enumerate(valid_frames):
new_prediction[i, frames] = prediction[i]
# The origin of (x, y) is in the upper right corner,
# while the (x,y) coordinates in the image are in the upper left corner.
distance = re_kpts[i, frames[1:], :, :2] - re_kpts[i, frames[:1], :, :2]
distance = np.mean(distance[:, [1, 4, 11, 14]], axis=-2, keepdims=True)
new_prediction[i, frames[1:], :, 0] -= distance[..., 0] / ratio_2d_3d
new_prediction[i, frames[1:], :, 1] += distance[..., 1] / ratio_2d_3d
# The origin of (x, y) is in the upper right corner,
# while the (x,y) coordinates in the image are in the upper left corner.
# Calculate the relative distance between two people
if len(valid_frames) == 2:
intersec_frames = [frame for frame in valid_frames[0] if frame in valid_frames[1]]
absolute_distance = re_kpts[0, intersec_frames[:1], :, :2] - re_kpts[1, intersec_frames[:1], :, :2]
absolute_distance = np.mean(absolute_distance[:, [1, 4, 11, 14]], axis=-2, keepdims=True) / 2.
new_prediction[0, valid_frames[0], :, 0] -= absolute_distance[..., 0] / ratio_2d_3d
new_prediction[0, valid_frames[0], :, 1] += absolute_distance[..., 1] / ratio_2d_3d
new_prediction[1, valid_frames[1], :, 0] += absolute_distance[..., 0] / ratio_2d_3d
new_prediction[1, valid_frames[1], :, 1] -= absolute_distance[..., 1] / ratio_2d_3d
# Pre-processing the case where the movement of Z axis is relatively large, such as 'sitting down'
# Remove the absolute distance
# new_prediction[:, :, 1:] -= new_prediction[:, :, :1]
# new_prediction[:, :, 0] = 0
new_prediction[:, :, :, 2] -= np.amin(new_prediction[:, :, :, 2])
return new_prediction
def revise_skes_real_time(prediction, re_kpts, width):
ratio_2d_3d_width = ratio_2d_3d * (width / 1920)
# prediction: (M, N, 3)
new_prediction = np.zeros((len(prediction), 17, 3), dtype=np.float32)
for i in range(len(prediction)):
new_prediction[i] = prediction[i]
initial_distance = re_kpts[i]
initial_distance = np.mean(initial_distance[[1, 4, 11, 14], :], axis=0)
new_prediction[i, :, 0] -= (initial_distance[0] - 3*width/5) / ratio_2d_3d_width
new_prediction[i, :, 1] += (initial_distance[1] - width/5) / ratio_2d_3d_width
new_prediction[:, :, 2] -= np.amin(new_prediction[:, :, 2])
return new_prediction
|