354 lines
11 KiB
Python
354 lines
11 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import logging
|
|
import time
|
|
|
|
import cv2
|
|
from argparse import ArgumentParser
|
|
|
|
from mmengine.logging import print_log
|
|
from mmpose.apis import inference_topdown, init_model
|
|
from mmpose.registry import VISUALIZERS
|
|
from mmpose.structures import merge_data_samples
|
|
import numpy as np
|
|
|
|
class Position:
|
|
def __init__(self, person_kpts):
|
|
if len(person_kpts) == 8:
|
|
self.left_elbow_angle = person_kpts[0]
|
|
self.right_elbow_angle = person_kpts[1]
|
|
self.left_shoulder_angle = person_kpts[2]
|
|
self.right_shoulder_angle = person_kpts[3]
|
|
self.left_leg_angle = person_kpts[4]
|
|
self.right_leg_angle = person_kpts[5]
|
|
self.left_hip_angle = person_kpts[6]
|
|
self.right_hip_angle = person_kpts[7]
|
|
else:
|
|
self.left_hip = person_kpts[11]
|
|
self.left_knee = person_kpts[13]
|
|
self.left_ankle = person_kpts[15]
|
|
|
|
self.right_hip = person_kpts[12]
|
|
self.right_knee = person_kpts[14]
|
|
self.right_ankle = person_kpts[16]
|
|
|
|
self.left_shoulder = person_kpts[5]
|
|
self.left_elbow = person_kpts[7]
|
|
self.left_wrist = person_kpts[9]
|
|
|
|
self.right_shoulder = person_kpts[6]
|
|
self.right_elbow = person_kpts[8]
|
|
self.right_wrist = person_kpts[10]
|
|
|
|
self.left_elbow_angle = angle_between(self.left_shoulder - self.left_elbow, self.left_wrist - self.left_elbow)
|
|
self.right_elbow_angle = angle_between(self.right_shoulder - self.right_elbow, self.right_wrist - self.right_elbow)
|
|
|
|
self.left_shoulder_angle = angle_between(self.right_shoulder - self.left_shoulder, self.left_elbow - self.left_shoulder)
|
|
self.right_shoulder_angle = angle_between(self.left_shoulder - self.right_shoulder, self.right_elbow - self.right_shoulder)
|
|
|
|
self.left_leg_angle = angle_between(self.left_hip - self.left_knee, self.left_ankle - self.left_knee)
|
|
self.right_leg_angle = angle_between(self.right_hip - self.right_knee, self.right_ankle - self.right_knee)
|
|
|
|
self.left_hip_angle = angle_between(self.left_shoulder - self.left_hip, self.left_knee - self.left_hip)
|
|
self.right_hip_angle = angle_between(self.right_shoulder - self.right_hip, self.right_knee - self.right_hip)
|
|
|
|
def to_array(self):
|
|
return [
|
|
self.left_elbow_angle,
|
|
self.right_elbow_angle,
|
|
self.left_shoulder_angle,
|
|
self.right_shoulder_angle,
|
|
self.left_leg_angle,
|
|
self.right_leg_angle,
|
|
self.left_hip_angle,
|
|
self.right_hip_angle
|
|
]
|
|
|
|
def distance_to(self, position):
|
|
error = 0
|
|
success = True
|
|
|
|
for i, x in enumerate(self.to_array()):
|
|
if 3 < i < 6:
|
|
continue
|
|
|
|
y = position.to_array()[i]
|
|
|
|
if abs(y) > 165:
|
|
x = abs(x)
|
|
y = abs(y)
|
|
|
|
dist = abs(y - x)
|
|
|
|
if dist > 20:
|
|
success = False
|
|
# print(f"{i} nie jest ok: moje: {x}, cel: {y}")
|
|
error += abs(y - x)
|
|
else:
|
|
pass
|
|
# print(f"{i} jest ok: moje: {x}, cel: {position.to_array()[i]}")
|
|
|
|
return (success, error)
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser()
|
|
# parser.add_argument('checkpoint', help='Checkpoint file')
|
|
parser.add_argument(
|
|
'--device', default='cuda:0', help='Device used for inference')
|
|
parser.add_argument(
|
|
'--draw-heatmap',
|
|
action='store_true',
|
|
help='Visualize the predicted heatmap')
|
|
parser.add_argument(
|
|
'--show-kpt-idx',
|
|
action='store_true',
|
|
default=False,
|
|
help='Whether to show the index of keypoints')
|
|
parser.add_argument(
|
|
'--skeleton-style',
|
|
default='mmpose',
|
|
type=str,
|
|
choices=['mmpose', 'openpose'],
|
|
help='Skeleton style selection')
|
|
parser.add_argument(
|
|
'--kpt-thr',
|
|
type=float,
|
|
default=0.3,
|
|
help='Visualizing keypoint thresholds')
|
|
parser.add_argument(
|
|
'--radius',
|
|
type=int,
|
|
default=3,
|
|
help='Keypoint radius for visualization')
|
|
parser.add_argument(
|
|
'--thickness',
|
|
type=int,
|
|
default=1,
|
|
help='Link thickness for visualization')
|
|
parser.add_argument(
|
|
'--alpha', type=float, default=0.8, help='The transparency of bboxes')
|
|
parser.add_argument(
|
|
'--camera-id', type=int, default=0, help='Camera device ID')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def angle_between(v1, v2):
|
|
"""
|
|
Liczy kąt między wektorami v1 i v2 w stopniach.
|
|
Znak kąta zależy od kierunku (przeciwnie do ruchu wskazówek zegara jest dodatni).
|
|
"""
|
|
# kąt w radianach
|
|
angle = np.arccos(np.clip(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)), -1.0, 1.0))
|
|
|
|
# obliczamy znak kąta w 2D: + jeśli v2 jest "po lewej" od v1
|
|
sign = np.sign(v1[0] * v2[1] - v1[1] * v2[0])
|
|
|
|
return np.degrees(angle) * sign
|
|
|
|
def is_visible(kpt, threshold=0.3):
|
|
# kpt = [x, y, score]
|
|
return kpt[2] > threshold
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
# build the model from a config file and a checkpoint file
|
|
if args.draw_heatmap:
|
|
cfg_options = dict(model=dict(test_cfg=dict(output_heatmaps=True)))
|
|
else:
|
|
cfg_options = None
|
|
|
|
model = init_model(
|
|
"mmpose/configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w48_8xb32-210e_coco-256x192.py",
|
|
"hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth",
|
|
device=args.device,
|
|
cfg_options=cfg_options)
|
|
|
|
# init visualizer
|
|
model.cfg.visualizer.radius = args.radius
|
|
model.cfg.visualizer.alpha = args.alpha
|
|
model.cfg.visualizer.line_width = args.thickness
|
|
|
|
visualizer = VISUALIZERS.build(model.cfg.visualizer)
|
|
visualizer.set_dataset_meta(
|
|
model.dataset_meta, skeleton_style=args.skeleton_style)
|
|
|
|
# start capturing video from camera
|
|
cap = cv2.VideoCapture(args.camera_id)
|
|
if not cap.isOpened():
|
|
print("Error: Cannot open camera")
|
|
return
|
|
|
|
lastDegs = []
|
|
lastMove = 0
|
|
visible = 0
|
|
maha = 0
|
|
start = False
|
|
|
|
oldpos = None
|
|
waiting = 100
|
|
|
|
positions = []
|
|
numPositions = 0
|
|
tPose = Position(
|
|
[
|
|
167.5568,
|
|
-161.67027,
|
|
-166.49443,
|
|
168.22028,
|
|
110.21745,
|
|
166.41733,
|
|
167.57822,
|
|
-176.08066
|
|
]
|
|
)
|
|
|
|
krakowiak = Position(
|
|
[np.float32(-98.226715), np.float32(106.39389), np.float32(-135.05656), np.float32(139.48904), np.float32(-149.9036), np.float32(8.216028), np.float32(174.70923), np.float32(-176.893)]
|
|
)
|
|
|
|
krakowiakRight = Position(
|
|
[np.float32(-110.61421), np.float32(-174.2424), np.float32(-127.05916), np.float32(174.9463),
|
|
np.float32(175.62007), np.float32(166.89127), np.float32(168.8219), np.float32(-178.02744)]
|
|
)
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("Error: Cannot read frame from camera")
|
|
break
|
|
|
|
# convert BGR to RGB
|
|
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# inference
|
|
batch_results = inference_topdown(model, img)
|
|
results = merge_data_samples(batch_results)
|
|
|
|
person_kpts = results.pred_instances.keypoints[0]
|
|
person_visible = results.pred_instances.keypoints_visible[0]
|
|
|
|
left_hand_visible = person_visible[5] > 0.75 and person_visible[5] > 0.75 and person_visible[9] > 0.75
|
|
right_hand_visible = person_visible[6] > 0.75 and person_visible[8] > 0.75 and person_visible[10] > 0.75
|
|
|
|
left_leg_visible = person_visible[11] > 0.75 and person_visible[13] > 0.75 and person_visible[15] > 0.75
|
|
right_leg_visible = person_visible[12] > 0.75 and person_visible[14] > 0.75 and person_visible[16] > 0.75
|
|
|
|
position = Position(person_kpts)
|
|
#
|
|
if position.distance_to(tPose)[0]:
|
|
print("\rT POSE!", end="")
|
|
elif position.distance_to(krakowiak)[0]:
|
|
print("\rKrakowiak", end="")
|
|
elif position.distance_to(krakowiakRight)[0]:
|
|
print("\rKrakowiak right", end="")
|
|
else:
|
|
print("\rNIC!", end="")
|
|
|
|
# print(position.left_elbow_angle)
|
|
# if oldpos is None:
|
|
# oldpos = position
|
|
#
|
|
# if oldpos.distance_to(position) < 7.5:
|
|
# print(f"\r{goalPosition.distance_to(position)}", end="")
|
|
#
|
|
# oldpos = position
|
|
|
|
# if error > 200:
|
|
|
|
# if start:
|
|
# if numPositions > 100:
|
|
# avgPosition = [0, 0, 0, 0, 0, 0, 0, 0]
|
|
#
|
|
# for element in positions:
|
|
# for i, data in enumerate(element):
|
|
# avgPosition[i] += data
|
|
#
|
|
# for i, element in enumerate(avgPosition):
|
|
# avgPosition[i] = avgPosition[i] / numPositions
|
|
#
|
|
# print(avgPosition)
|
|
#
|
|
# break
|
|
# else:
|
|
# print(f"\r{numPositions}", end="")
|
|
# if oldpos is None:
|
|
# oldpos = position
|
|
#
|
|
# if oldpos.distance_to(position)[0]:
|
|
# positions.insert(0, position.to_array())
|
|
# numPositions += 1
|
|
#
|
|
# oldpos = position
|
|
#
|
|
# # else:
|
|
# # print(f"\rOK!", end="")
|
|
#
|
|
# if waiting != 0 and not start:
|
|
# print(f"\r{waiting}", end="")
|
|
# waiting -= 1
|
|
# else:
|
|
# start = True
|
|
|
|
# lastDegs.insert(0, (left_elbow_angle, right_elbow_angle))
|
|
#
|
|
# last = 0
|
|
# lastCount = 0
|
|
#
|
|
# for element in lastDegs:
|
|
# last += element[1]
|
|
# lastCount += 1
|
|
#
|
|
# last = last / lastCount
|
|
# dist = right_elbow_angle - last
|
|
#
|
|
# if not right_visible:
|
|
# visible = 0
|
|
# print("\rNie widać prawej ręki!!!!!!!!", end="")
|
|
# else:
|
|
# if maha == 0:
|
|
# print("\rWidać rękę, nie maha!", end="")
|
|
# else:
|
|
# maha -= 1
|
|
#
|
|
# visible += 1
|
|
#
|
|
# if 15 < abs(dist) < 60 and visible > 5:
|
|
# if lastMove != dist > 0:
|
|
# maha = 10
|
|
# print("\rmaha!", end="")
|
|
#
|
|
# lastMove = dist > 0
|
|
#
|
|
# if len(lastDegs) > 5:
|
|
# lastDegs.pop()
|
|
|
|
# visualize
|
|
vis_img = visualizer.add_datasample(
|
|
'result',
|
|
img,
|
|
data_sample=results,
|
|
draw_gt=False,
|
|
draw_bbox=True,
|
|
kpt_thr=args.kpt_thr,
|
|
draw_heatmap=args.draw_heatmap,
|
|
show_kpt_idx=args.show_kpt_idx,
|
|
skeleton_style=args.skeleton_style,
|
|
show=False,
|
|
out_file=None)
|
|
|
|
# convert RGB back to BGR for OpenCV
|
|
vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
|
|
cv2.imshow('Live Pose Estimation', vis_img)
|
|
|
|
# press 'q' to quit
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|