182 lines
6.9 KiB
Python
182 lines
6.9 KiB
Python
import time
|
|
|
|
import cv2
|
|
import math
|
|
|
|
from MNPoseDetection.consensus import *
|
|
from MNPoseDetection.pose_detection import PoseDetection
|
|
|
|
# 定义人体骨架连接
|
|
skeleton = [
|
|
('nose', 'left_eye'), ('nose', 'right_eye'), ('left_eye', 'left_ear'), ('right_eye', 'right_ear'),
|
|
('left_shoulder', 'right_shoulder'), ('left_shoulder', 'left_elbow'),
|
|
('right_shoulder', 'right_elbow'), ('left_elbow', 'left_wrist'), ('right_elbow', 'right_wrist'),
|
|
('left_shoulder', 'left_hip'), ('right_shoulder', 'right_hip'), ('left_hip', 'right_hip'),
|
|
('left_hip', 'left_knee'), ('right_hip', 'right_knee'), ('left_knee', 'left_ankle'), ('right_knee', 'right_ankle')
|
|
]
|
|
|
|
|
|
class MoveNetAlgorithmPlugin:
|
|
detect_corner = None
|
|
config = (0.6, 50, 55, 60)
|
|
pose = None
|
|
|
|
@classmethod
|
|
def class_init(cls):
|
|
MoveNetAlgorithmPlugin.pose = PoseDetection(*MoveNetAlgorithmPlugin.config)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.lm_list = []
|
|
if MoveNetAlgorithmPlugin.pose is None:
|
|
MoveNetAlgorithmPlugin.class_init()
|
|
|
|
@classmethod
|
|
def get_config(cls):
|
|
return MoveNetAlgorithmPlugin.config
|
|
|
|
@classmethod
|
|
def set_config(cls, config):
|
|
if config != MoveNetAlgorithmPlugin.config:
|
|
new_config = list(config)
|
|
MoveNetAlgorithmPlugin.config = new_config
|
|
MoveNetAlgorithmPlugin.pose = PoseDetection(*MoveNetAlgorithmPlugin.config)
|
|
|
|
@classmethod
|
|
def get_corner(cls):
|
|
return MoveNetAlgorithmPlugin.detect_corner
|
|
|
|
@classmethod
|
|
def set_corner(cls, corner):
|
|
if corner != MoveNetAlgorithmPlugin.detect_corner:
|
|
MoveNetAlgorithmPlugin.detect_corner = corner
|
|
|
|
def set_result(self, lm_list):
|
|
self.lm_list = lm_list
|
|
|
|
@classmethod
|
|
def find_pose(cls, img):
|
|
dc = MoveNetAlgorithmPlugin.detect_corner
|
|
if dc:
|
|
img_cropped = img[dc[0]:dc[1], dc[2]:dc[3]]
|
|
else:
|
|
img_cropped = img
|
|
img_rgb = cv2.cvtColor(img_cropped, cv2.COLOR_BGR2RGB)
|
|
results = MoveNetAlgorithmPlugin.pose.detect(img_rgb)
|
|
return results, img_cropped
|
|
|
|
@staticmethod
|
|
def find_pose_with_drawing(img, draw=True):
|
|
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
results = MoveNetAlgorithmPlugin.pose.detect(img_rgb)
|
|
if draw:
|
|
for part_a, part_b in skeleton:
|
|
x_a, y_a = results[part_a][KEY_POINTS]
|
|
x_b, y_b = results[part_b][KEY_POINTS]
|
|
if (
|
|
results[part_a][RELIABLE]
|
|
and results[part_b][RELIABLE]
|
|
and results[part_a][PASS_TEST]
|
|
and results[part_b][PASS_TEST]
|
|
):
|
|
cv2.line(img, (x_a, y_a), (x_b, y_b), (255, 64, 64), 2)
|
|
else:
|
|
cv2.line(img, (x_a, y_a), (x_b, y_b), (125, 125, 255), 2)
|
|
for label, info in results.items():
|
|
keypoint = info[KEY_POINTS]
|
|
reliable = info[RELIABLE]
|
|
pass_test = info[PASS_TEST]
|
|
if not reliable or not pass_test:
|
|
cv2.circle(img, keypoint, 4, (0, 64, 255), -1)
|
|
else:
|
|
cv2.circle(img, keypoint, 4, (0, 255, 125), -1)
|
|
return img
|
|
|
|
def drawPoint(self, img, p1, p2, p3, bias_x=0, bias_y=0):
|
|
if len(self.lm_list) == 0:
|
|
return
|
|
x1, y1 = self.lm_list[p1][KEY_POINTS][0], self.lm_list[p1][KEY_POINTS][1]
|
|
x2, y2 = self.lm_list[p2][KEY_POINTS][0], self.lm_list[p2][KEY_POINTS][1]
|
|
x3, y3 = self.lm_list[p3][KEY_POINTS][0], self.lm_list[p3][KEY_POINTS][1]
|
|
# print((x1, y1), (x2, y2), (x3, y3))
|
|
cv2.circle(img, (x1 + bias_x, y1 + bias_y), 5, (102, 106, 233), cv2.FILLED)
|
|
cv2.circle(img, (x2 + bias_x, y2 + bias_y), 5, (50, 80, 4), cv2.FILLED)
|
|
cv2.circle(img, (x3 + bias_x, y3 + bias_y), 5, (20, 95, 104), cv2.FILLED)
|
|
|
|
def drawPoint_more(self, img, queue, bias_x=0, bias_y=0):
|
|
if len(self.lm_list) == 0:
|
|
return
|
|
|
|
color_list = [(0, 255, 255), (255, 255, 0), (255, 0, 255), (0.0, 255), (0, 255, 0), (255, 255, 255)]
|
|
|
|
xy_list = []
|
|
for i in queue:
|
|
x, y = self.lm_list[i][KEY_POINTS][0], self.lm_list[i][KEY_POINTS][1]
|
|
xy_list.append([x, y])
|
|
|
|
for i in range(len(xy_list)):
|
|
x = xy_list[i][0]
|
|
y = xy_list[i][1]
|
|
c = i % 6
|
|
cv2.circle(img, (x + bias_x, y + bias_y), 5, color_list[c], cv2.FILLED)
|
|
|
|
def findPosition(self, img, draw=True):
|
|
if self.lm_list:
|
|
if draw:
|
|
for label, result in self.lm_list:
|
|
cx, cy = result[KEY_POINTS]
|
|
cv2.circle(img, (cx, cy), 5, (255, 0, 0), cv2.FILLED)
|
|
return self.lm_list
|
|
|
|
def findConcave(self, p1, p2, p3):
|
|
x1, y1 = self.lm_list[p1][KEY_POINTS][0], self.lm_list[p1][KEY_POINTS][1]
|
|
x2, y2 = self.lm_list[p2][KEY_POINTS][0], self.lm_list[p2][KEY_POINTS][1]
|
|
x3, y3 = self.lm_list[p3][KEY_POINTS][0], self.lm_list[p3][KEY_POINTS][1]
|
|
|
|
# Calculate Angle
|
|
Concave_angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2))
|
|
if Concave_angle < 0:
|
|
Concave_angle += 360
|
|
|
|
return Concave_angle
|
|
|
|
def findAngle(self, img, p1, p2, p3, draw=True):
|
|
# Get the landmarks
|
|
x1, y1 = self.lm_list[p1][KEY_POINTS][0], self.lm_list[p1][KEY_POINTS][1]
|
|
x2, y2 = self.lm_list[p2][KEY_POINTS][0], self.lm_list[p2][KEY_POINTS][1]
|
|
x3, y3 = self.lm_list[p3][KEY_POINTS][0], self.lm_list[p3][KEY_POINTS][1]
|
|
|
|
# Calculate Angle
|
|
angle = math.degrees(math.atan2(y3 - y2, x3 - x2) -
|
|
math.atan2(y1 - y2, x1 - x2))
|
|
if angle < 0:
|
|
angle += 360
|
|
if angle > 180:
|
|
angle = 360 - angle
|
|
elif angle > 180:
|
|
angle = 360 - angle
|
|
|
|
# Draw
|
|
if draw:
|
|
cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 2)
|
|
cv2.line(img, (x3, y3), (x2, y2), (255, 255, 255), 2)
|
|
|
|
cv2.circle(img, (x1, y1), 3, (0, 0, 255), cv2.FILLED)
|
|
cv2.circle(img, (x1, y1), 10, (0, 0, 255), 2)
|
|
cv2.circle(img, (x2, y2), 3, (0, 0, 255), cv2.FILLED)
|
|
cv2.circle(img, (x2, y2), 10, (0, 0, 255), 2)
|
|
cv2.circle(img, (x3, y3), 3, (0, 0, 255), cv2.FILLED)
|
|
cv2.circle(img, (x3, y3), 10, (0, 0, 255), 2)
|
|
|
|
cv2.putText(img, str(int(angle)), (x2 - 20, y2 + 20),
|
|
cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)
|
|
return angle
|
|
|
|
def findIncludedAngle(self, x1, y1, x2, y2):
|
|
IncludedAngle = math.asin((y1 - y2) / math.hypot((y1 - y2), (x1 - x2))) * 180 / math.pi
|
|
return IncludedAngle
|
|
|
|
def findRange(self, x1, y1, x2, y2):
|
|
Range = math.sqrt(math.pow(x1 - x2, 2) + math.pow(y1 - y2, 2))
|
|
return Range
|