220 lines
9.1 KiB
Python
220 lines
9.1 KiB
Python
import cv2
|
|
import mediapipe as mp
|
|
import numpy as np
|
|
|
|
mpHands = mp.solutions.hands
|
|
mpPose = mp.solutions.pose
|
|
mpDraw = mp.solutions.drawing_utils
|
|
mpHolistic = mp.solutions.holistic
|
|
|
|
import logging
|
|
FORMAT = '%(asctime)-15s %(message)s'
|
|
logging.basicConfig(filename="hpe_mp_class.log", level=logging.INFO, format=FORMAT)
|
|
logger = logging.getLogger("hpe_mp_class_logger")
|
|
|
|
from ModelUE4 import bodyconvert
|
|
from ModelUE4 import rhandconvert
|
|
from ModelUE4 import lhandconvert
|
|
from ModelUE4 import bodyaposeworld
|
|
from ModelUE4 import bodyaposelocal
|
|
from ModelUE4 import bodyconvertlocal
|
|
|
|
class hpe_mp_class():
|
|
|
|
def __init__(self):
|
|
hands_static_image_mode = False
|
|
hands_max_num_hands = 2
|
|
hands_min_detection_confidence = 0.7
|
|
hands_min_tracking_confidence = 0.5
|
|
|
|
pose_static_image_mode = False
|
|
pose_upper_body_only = False
|
|
pose_smooth_landmarks = True
|
|
pose_min_detection_confidence = 0.7
|
|
pose_min_tracking_confidence = 0.5
|
|
|
|
hol_static_image_mode = False
|
|
hol_upper_body_only = False
|
|
hol_smooth_landmarks = True
|
|
hol_min_detection_confidence = 0.7
|
|
hol_min_tracking_confidence = 0.5
|
|
|
|
self.holistic_use = True
|
|
self.coef = 1.0
|
|
|
|
try:
|
|
self.hands = mpHands.Hands(min_detection_confidence=hands_min_detection_confidence,
|
|
min_tracking_confidence=hands_min_tracking_confidence)
|
|
self.poses = mpPose.Pose(min_detection_confidence=pose_min_detection_confidence,
|
|
min_tracking_confidence=pose_min_tracking_confidence)
|
|
self.hol = mpHolistic.Holistic(min_detection_confidence=hol_min_detection_confidence,
|
|
min_tracking_confidence=hol_min_tracking_confidence)
|
|
logger.info("Success initialize hpe class")
|
|
except Exception as err:
|
|
logger.exception("Error initialize hpe class: " + str(err))
|
|
|
|
def process(self, image, scale_pose=0.42):
|
|
try:
|
|
imgRGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
if self.holistic_use:
|
|
self.results_hol = self.hol.process(image)
|
|
|
|
poseslms = {}
|
|
if self.results_hol.pose_landmarks:
|
|
for id, lm in enumerate(self.results_hol.pose_landmarks.landmark):
|
|
poseslms[id] = lm
|
|
rast = np.sqrt(pow(poseslms[11].x - poseslms[12].x, 2) + pow(poseslms[11].y - poseslms[12].y, 2) + pow(poseslms[11].z - poseslms[12].z, 2))
|
|
self.coef = scale_pose / rast
|
|
else:
|
|
self.results_hands = self.hands.process(imgRGB)
|
|
self.results_pose = self.poses.process(imgRGB)
|
|
except Exception as err:
|
|
logger.exception("Error processing hpe class: " + str(err))
|
|
|
|
def show(self, image):
|
|
try:
|
|
if self.holistic_use:
|
|
mpDraw.draw_landmarks(
|
|
image,
|
|
self.results_hol.face_landmarks,
|
|
mpHolistic.FACE_CONNECTIONS)
|
|
mpDraw.draw_landmarks(
|
|
image,
|
|
self.results_hol.pose_landmarks,
|
|
mpHolistic.POSE_CONNECTIONS)
|
|
mpDraw.draw_landmarks(
|
|
image,
|
|
self.results_hol.left_hand_landmarks,
|
|
mpHands.HAND_CONNECTIONS)
|
|
mpDraw.draw_landmarks(
|
|
image,
|
|
self.results_hol.right_hand_landmarks,
|
|
mpHands.HAND_CONNECTIONS)
|
|
else:
|
|
if self.results_hands.multi_hand_landmarks:
|
|
for handLms in self.results_hands.multi_hand_landmarks:
|
|
mpDraw.draw_landmarks(image, handLms, mpHands.HAND_CONNECTIONS)
|
|
for id, lm in enumerate(handLms.landmark):
|
|
h, w, c = image.shape
|
|
cx, cy = int(lm.x * w), int(lm.y * h)
|
|
if id == 4 or id == 8 or id == 12 or id == 16 or id == 20:
|
|
cv2.circle(image, (cx, cy), 10, (255, 0, 0), cv2.FILLED)
|
|
|
|
if self.results_pose.pose_landmarks:
|
|
mpDraw.draw_landmarks(image, self.results_pose.pose_landmarks, mpPose.POSE_CONNECTIONS)
|
|
except Exception as err:
|
|
logger.exception("Error showing hpe class: " + str(err))
|
|
|
|
def getResults(self):
|
|
res = {}
|
|
try:
|
|
poseslms = {}
|
|
handslms = []
|
|
|
|
if self.holistic_use:
|
|
if self.results_hol.pose_landmarks:
|
|
for id, lm in enumerate(self.results_hol.pose_landmarks.landmark):
|
|
poseslms[id] = lm
|
|
hand = {}
|
|
if self.results_hol.left_hand_landmarks:
|
|
for id, lm in enumerate(self.results_hol.left_hand_landmarks.landmark):
|
|
hand[id] = lm
|
|
handslms.append(hand)
|
|
if self.results_hol.right_hand_landmarks:
|
|
for id, lm in enumerate(self.results_hol.right_hand_landmarks.landmark):
|
|
hand[id+21] = lm
|
|
handslms.append(hand)
|
|
res["poses"] = poseslms
|
|
res["hands"] = handslms
|
|
else:
|
|
if self.results_pose.pose_landmarks:
|
|
for id, lm in enumerate(self.results_pose.pose_landmarks.landmark):
|
|
poseslms[id] = lm
|
|
if self.results_hands.multi_hand_landmarks:
|
|
for handLms in self.results_hands.multi_hand_landmarks:
|
|
hand = {}
|
|
for id, lm in enumerate(handLms.landmark):
|
|
hand[id] = lm
|
|
handslms.append(hand)
|
|
res["poses"] = poseslms
|
|
res["hands"] = handslms
|
|
except Exception as err:
|
|
logger.exception("Error getting result hpe class: " + str(err))
|
|
|
|
return res
|
|
|
|
def scaleResult(self, result):
|
|
try:
|
|
if self.holistic_use:
|
|
for pp in result["poses"].keys():
|
|
result["poses"][pp].x *= self.coef
|
|
result["poses"][pp].y *= self.coef
|
|
result["poses"][pp].z *= self.coef
|
|
|
|
if len(result["hands"]) > 0:
|
|
for hp0 in result["hands"][0].keys():
|
|
result["hands"][0][hp0].x *= self.coef
|
|
result["hands"][0][hp0].y *= self.coef
|
|
result["hands"][0][hp0].z *= self.coef
|
|
|
|
if len(result["hands"]) > 1:
|
|
for hp1 in result["hands"][1].keys():
|
|
result["hands"][1][hp1].x *= self.coef
|
|
result["hands"][1][hp1].y *= self.coef
|
|
result["hands"][1][hp1].z *= self.coef
|
|
except Exception as err:
|
|
logger.exception("Error scaling hpe class: " + str(err))
|
|
|
|
def getJSON(self, apose=False, world=True):
|
|
data = {}
|
|
try:
|
|
if apose:
|
|
if world:
|
|
bodyaposeworld(data)
|
|
else:
|
|
bodyaposelocal(data)
|
|
else:
|
|
if world:
|
|
if self.holistic_use:
|
|
poseslms = {}
|
|
maxy = 0
|
|
if self.results_hol.pose_landmarks:
|
|
for id, lm in enumerate(self.results_hol.pose_landmarks.landmark):
|
|
poseslms[id] = lm
|
|
if lm.y > maxy:
|
|
maxy = lm.y
|
|
|
|
bodyconvert(poseslms, data, self.coef, maxy)
|
|
|
|
rhandlms = {}
|
|
if self.results_hol.right_hand_landmarks:
|
|
for id, lm in enumerate(self.results_hol.right_hand_landmarks.landmark):
|
|
rhandlms[id] = lm
|
|
|
|
rhandconvert(rhandlms, data, self.coef)
|
|
|
|
lhandlms = {}
|
|
if self.results_hol.left_hand_landmarks:
|
|
for id, lm in enumerate(self.results_hol.left_hand_landmarks.landmark):
|
|
lhandlms[id] = lm
|
|
|
|
lhandconvert(lhandlms, data, self.coef)
|
|
else:
|
|
bodyaposelocal(data)
|
|
if self.holistic_use:
|
|
poseslms = {}
|
|
maxy = 0
|
|
if self.results_hol.pose_landmarks:
|
|
for id, lm in enumerate(self.results_hol.pose_landmarks.landmark):
|
|
poseslms[id] = lm
|
|
if lm.y > maxy:
|
|
maxy = lm.y
|
|
|
|
bodyconvert(poseslms, data, self.coef, maxy)
|
|
|
|
bodyconvertlocal(poseslms, data, self.coef, maxy)
|
|
except Exception as err:
|
|
logger.exception("Error json converting hpe class: " + str(err))
|
|
|
|
return data |