bone_helper.py
import json
import cv2
import requests
def add_frame(json_data, cap, out, video_save):
frame_count = 0
a1 = []
b1 = []
t = len(json_data["skeleton_info"][0]["skeleton_coordinate"])
a2 = []
b2 = []
c2 = []
c1 = []
while 1:
a1.clear()
b1.clear()
a2.clear()
b2.clear()
c1.clear()
c2.clear()
ret, frame = cap.read()
for i in range(25):
a1.append(
int(json_data["skeleton_info"][frame_count]["skeleton_coordinate"][0]["pose_keypoints_2d"][i * 2]))
b1.append(
int(json_data["skeleton_info"][frame_count]["skeleton_coordinate"][0]["pose_keypoints_2d"][
i * 2 + 1]))
if t == 2:
a2.append(
int(json_data["skeleton_info"][frame_count]["skeleton_coordinate"][1]["pose_keypoints_2d"][
i * 2]))
b2.append(
int(json_data["skeleton_info"][frame_count]["skeleton_coordinate"][1]["pose_keypoints_2d"][
i * 2 + 1]))
if a1[i] != 0 | b1[i] != 0:
c1.append(1)
cv2.circle(frame, (a1[i], b1[i]), 5, (255, 0, 0), -1)
else:
c1.append(0)
if t == 2:
if a2[i] != 0 | b2[i] != 0:
c2.append(1)
cv2.circle(frame, (a2[i], b2[i]), 5, (255, 0, 0), -1)
else:
c2.append(0)
if c1[0] and c1[1]:
cv2.line(frame, (a1[0], b1[0]), (a1[1], b1[1]), (255, 0, 0), 3)
if c1[1] and c1[2]:
cv2.line(frame, (a1[1], b1[1]), (a1[2], b1[2]), (255, 0, 0), 3)
if c1[2] and c1[3]:
cv2.line(frame, (a1[2], b1[2]), (a1[3], b1[3]), (255, 0, 0), 3)
if c1[3] and c1[4]:
cv2.line(frame, (a1[3], b1[3]), (a1[4], b1[4]), (255, 0, 0), 3)
if c1[1] and c1[5]:
cv2.line(frame, (a1[1], b1[1]), (a1[5], b1[5]), (255, 0, 0), 3)
if c1[5] and c1[6]:
cv2.line(frame, (a1[5], b1[5]), (a1[6], b1[6]), (255, 0, 0), 3)
if c1[6] and c1[7]:
cv2.line(frame, (a1[6], b1[6]), (a1[7], b1[7]), (255, 0, 0), 3)
if c1[1] and c1[8]:
cv2.line(frame, (a1[1], b1[1]), (a1[8], b1[8]), (255, 0, 0), 3)
if c1[8] and c1[9]:
cv2.line(frame, (a1[8], b1[8]), (a1[9], b1[9]), (255, 0, 0), 3)
if c1[9] and c1[10]:
cv2.line(frame, (a1[9], b1[9]), (a1[10], b1[10]), (255, 0, 0), 3)
if c1[10] and c1[11]:
cv2.line(frame, (a1[10], b1[10]), (a1[11], b1[11]), (255, 0, 0), 3)
if c1[8] and c1[12]:
cv2.line(frame, (a1[8], b1[8]), (a1[12], b1[12]), (255, 0, 0), 3)
if c1[12] and c1[13]:
cv2.line(frame, (a1[12], b1[12]), (a1[13], b1[13]), (255, 0, 0), 3)
if c1[13] and c1[14]:
cv2.line(frame, (a1[13], b1[13]), (a1[14], b1[14]), (255, 0, 0), 3)
if c1[0] and c1[15]:
cv2.line(frame, (a1[0], b1[0]), (a1[15], b1[15]), (255, 0, 0), 3)
if c1[0] and c1[16]:
cv2.line(frame, (a1[0], b1[0]), (a1[16], b1[16]), (255, 0, 0), 3)
if c1[15] and c1[17]:
cv2.line(frame, (a1[15], b1[15]), (a1[17], b1[17]), (255, 0, 0), 3)
if c1[16] and c1[18]:
cv2.line(frame, (a1[16], b1[16]), (a1[18], b1[18]), (255, 0, 0), 3)
if c1[14] and c1[21]:
cv2.line(frame, (a1[14], b1[14]), (a1[21], b1[21]), (255, 0, 0), 3)
if c1[14] and c1[19]:
cv2.line(frame, (a1[14], b1[14]), (a1[19], b1[19]), (255, 0, 0), 3)
if c1[19] and c1[20]:
cv2.line(frame, (a1[19], b1[19]), (a1[20], b1[20]), (255, 0, 0), 3)
if c1[11] and c1[24]:
cv2.line(frame, (a1[11], b1[11]), (a1[24], b1[24]), (255, 0, 0), 3)
if c1[11] and c1[22]:
cv2.line(frame, (a1[11], b1[11]), (a1[22], b1[22]), (255, 0, 0), 3)
if c1[22] and c1[23]:
cv2.line(frame, (a1[22], b1[22]), (a1[23], b1[23]), (255, 0, 0), 3)
if t == 2:
if c2[0] and c2[1]:
cv2.line(frame, (a2[0], b2[0]), (a2[1], b2[1]), (255, 0, 0), 3)
if c2[1] and c2[2]:
cv2.line(frame, (a2[1], b2[1]), (a2[2], b2[2]), (255, 0, 0), 3)
if c2[2] and c2[3]:
cv2.line(frame, (a2[2], b2[2]), (a2[3], b2[3]), (255, 0, 0), 3)
if c2[3] and c2[4]:
cv2.line(frame, (a2[3], b2[3]), (a2[4], b2[4]), (255, 0, 0), 3)
if c2[1] and c2[5]:
cv2.line(frame, (a2[1], b2[1]), (a2[5], b2[5]), (255, 0, 0), 3)
if c2[5] and c2[6]:
cv2.line(frame, (a2[5], b2[5]), (a2[6], b2[6]), (255, 0, 0), 3)
if c2[6] and c2[7]:
cv2.line(frame, (a2[6], b2[6]), (a2[7], b2[7]), (255, 0, 0), 3)
if c2[1] and c2[8]:
cv2.line(frame, (a2[1], b2[1]), (a2[8], b2[8]), (255, 0, 0), 3)
if c2[8] and c2[9]:
cv2.line(frame, (a2[8], b2[8]), (a2[9], b2[9]), (255, 0, 0), 3)
if c2[9] and c2[10]:
cv2.line(frame, (a2[9], b2[9]), (a2[10], b2[10]), (255, 0, 0), 3)
if c2[10] and c2[11]:
cv2.line(frame, (a2[10], b2[10]), (a2[11], b2[11]), (255, 0, 0), 3)
if c2[8] and c2[12]:
cv2.line(frame, (a2[8], b2[8]), (a2[12], b2[12]), (255, 0, 0), 3)
if c2[12] and c2[13]:
cv2.line(frame, (a2[12], b2[12]), (a2[13], b2[13]), (255, 0, 0), 3)
if c2[13] and c2[14]:
cv2.line(frame, (a2[13], b2[13]), (a2[14], b2[14]), (255, 0, 0), 3)
if c2[0] and c2[15]:
cv2.line(frame, (a2[0], b2[0]), (a2[15], b2[15]), (255, 0, 0), 3)
if c2[0] and c2[16]:
cv2.line(frame, (a2[0], b2[0]), (a2[16], b2[16]), (255, 0, 0), 3)
if c2[15] and c2[17]:
cv2.line(frame, (a2[15], b2[15]), (a2[17], b2[17]), (255, 0, 0), 3)
if c2[16] and c2[18]:
cv2.line(frame, (a2[16], b2[16]), (a2[18], b2[18]), (255, 0, 0), 3)
if c2[14] and c2[21]:
cv2.line(frame, (a2[14], b2[14]), (a2[21], b2[21]), (255, 0, 0), 3)
if c2[14] and c2[19]:
cv2.line(frame, (a2[14], b2[14]), (a2[19], b2[19]), (255, 0, 0), 3)
if c2[19] and c2[20]:
cv2.line(frame, (a2[19], b2[19]), (a2[20], b2[20]), (255, 0, 0), 3)
if c2[11] and c2[24]:
cv2.line(frame, (a2[11], b2[11]), (a2[24], b2[24]), (255, 0, 0), 3)
if c2[11] and c2[22]:
cv2.line(frame, (a2[11], b2[11]), (a2[22], b2[22]), (255, 0, 0), 3)
if c2[22] and c2[23]:
cv2.line(frame, (a2[22], b2[22]), (a2[23], b2[23]), (255, 0, 0), 3)
if frame_count < len(json_data["skeleton_info"]) - 1:
frame_count += 1
else:
break
if video_save:
out.write(frame)
return video_save
def draw_bone(json_data_single, video_src):
try:
video_save = str(video_src).replace(".mp4", "_b.mp4")
cap = cv2.VideoCapture(video_src)
frame_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
frame_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
fps = cap.get(cv2.CAP_PROP_FPS)
if video_save:
fourcc = cv2.VideoWriter_fourcc(*'avc1')
out = cv2.VideoWriter(video_save, fourcc, fps, (int(frame_width), int(frame_height)))
add_frame(json_data=json_data_single, cap=cap, out=out, video_save=str(video_save))
cap.release()
cv2.destroyWindow('cap')
return video_save
except Exception as e:
print(e)
def draw_bones(video_src, json_path_list, callback_url):
video_save_path = BoneHelper(video_src=video_src,
json_path_list=json_path_list).draw_bones()
if video_save_path:
requests.post(callback_url, data=json.dumps({
'succeed': True,
'videoPath': video_save_path
}))
else:
requests.post(callback_url, data=json.dumps({
'succeed': False,
'videoPath': video_save_path
}))
return video_save_path
class BoneHelper:
def __init__(self, json_path_list, video_src):
self.json_path_list = json_path_list
self.video_src = video_src
def read_json_list(self):
list_json_data = []
for json_path in self.json_path_list:
if json_path is None:
print("need more info")
continue
else:
with open(json_path, 'r', encoding='utf8')as fp:
json_data = json.load(fp)
list_json_data.append(json_data)
return list_json_data
def draw_bones(self):
json_list = self.read_json_list()
if json_list:
try:
video_src = self.video_src
for json_data in json_list:
video_src = draw_bone(json_data_single=json_data,
video_src=video_src)
return video_src
except Exception as e:
print(e)
return None
if __name__ == '__main__':
draw_bones(video_src=r'G:\object_track\video\20210110165018-4.mp4', json_path_list=[
r'G:\object_track\video\20210110165018-3_right.mp4.json',
r'G:\object_track\video\20210110165018-3_left.mp4.json',
])
对外接口
@app.route('/add/bones/', methods=['post'])
def add_bone():
try:
data = json.loads(request.data)
video_src = str(data["video_dir"])
video_json_list = str(data["video_json_list"])
callback_url = str(data["callback_url"])
executor.submit(draw_bones(), video_json_list, video_src, callback_url)
return json.dumps({
'succeed': True,
'msg': video_json_path + 'is processing'
}).encode(encoding="utf-8")
except Exception as e:
return json.dumps({
'succeed': False,
'msg': e
}).encode(encoding="utf-8")
