eBIM-Satu2 / utils.py
chai991014's picture
beta version
2097656
Raw
History Blame Contribute Delete
3.22 kB
import cv2
import numpy as np
import mediapipe as mp
mp_holistic = mp.solutions.holistic # Holistic model
mp_drawing = mp.solutions.drawing_utils # Drawing utilities
def mediapipe_detection(image, model):
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Color conversion from BGR to RGB
image.flags.writeable = False # Image is no longer writeable
results = model.process(image) # Make prediction
image.flags.writeable = True # Image is no longer writeable
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # Color conversion RGB to BGR
return image, results
def draw_styled_landmarks(image,results):
# Draw pose connection
mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=1, circle_radius=1),
mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1)
)
# Draw left hand connection
mp_drawing.draw_landmarks(image, results.left_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=1, circle_radius=2),
mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1)
)
# Draw right hand connection
mp_drawing.draw_landmarks(image, results.right_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=1, circle_radius=2),
mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1)
)
def extract_keypoints(results):
pose = np.array([[res.x, res.y, res.z, res.visibility] for res in results.pose_landmarks.landmark]).flatten() if results.pose_landmarks else np.zeros(33*4)
lh = np.array([[res.x, res.y, res.z] for res in results.left_hand_landmarks.landmark]).flatten() if results.left_hand_landmarks else np.zeros(21*3)
rh = np.array([[res.x, res.y, res.z] for res in results.right_hand_landmarks.landmark]).flatten() if results.right_hand_landmarks else np.zeros(21*3)
return np.concatenate([pose, lh, rh])
# --- STGCN Helpers ---
def get_adjacency_matrix():
A = np.eye(75)
for conn in mp_holistic.POSE_CONNECTIONS:
A[conn[0], conn[1]] = 1;
A[conn[1], conn[0]] = 1
for conn in mp_holistic.HAND_CONNECTIONS:
A[conn[0] + 33, conn[1] + 33] = 1;
A[conn[1] + 33, conn[0] + 33] = 1
A[conn[0] + 54, conn[1] + 54] = 1;
A[conn[1] + 54, conn[0] + 54] = 1
return A
def reshape_for_stgcn(X):
N, T, _ = X.shape
X_new = np.zeros((N, T, 75, 3))
for i in range(N):
for t in range(T):
frame = X[i, t]
pose = frame[0:132].reshape(33, 4)[:, :3]
lh = frame[132:195].reshape(21, 3)
rh = frame[195:258].reshape(21, 3)
X_new[i, t] = np.concatenate([pose, lh, rh], axis=0)
return X_new.transpose(0, 3, 1, 2)