import pandas as pd import numpy as np import cv2 def extract_landmarks(image, mp_pose, cols): pre_list = [] with mp_pose.Pose(static_image_mode=True, enable_segmentation=True) as pose: result = pose.process( cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) try: # xy = bounding_box(result.pose_landmarks.landmark) for landmark in result.pose_landmarks.landmark: pre_list.append(landmark) predict = True except AttributeError: return True, pd.DataFrame(), None if predict == True: gen1116 = np.array([ [ pre_list[m].x, pre_list[m].y, pre_list[m].z, pre_list[m].visibility ] for m in range(11, 17) ]).flatten().tolist() gen2333 = np.array([ [ pre_list[m].x, pre_list[m].y, pre_list[m].z, pre_list[m].visibility ] for m in range(23, 33) ]).flatten().tolist() gen1116.extend(gen2333) all_list = [ pre_list[0].x, pre_list[0].y, pre_list[0].z, pre_list[0].visibility, ] all_list.extend(gen1116) return False, pd.DataFrame([all_list], columns=cols), result.pose_landmarks # def bounding_box(landmarks): # w = 1280 # h = 720 # xy = [0, 0, w, h] # for landmark in landmarks: # x, y = int(landmark.x * w), int(landmark.y * h) # if x > xy[0]: # xy[0] = x # if x < xy[2]: # xy[2] = x # if y > xy[1]: # xy[1] = y # if y < xy[3]: # xy[3] = y # return xy