spoter_mod/skeleton_extractor.py

import warnings

import pandas as pd
from os import path
import cv2
import mediapipe as mp
import json
from spoter_mod.pose_model_identifier import BODY_IDENTIFIERS, HAND_IDENTIFIERS, mp_holistic_data

mp_drawing = mp.solutions.drawing_utils
mp_holistic = mp.solutions.holistic
mp_drawing_styles = mp.solutions.drawing_styles

holistic = mp_holistic.Holistic()

column_names = []
column_names.append('video_id')
for id_name in BODY_IDENTIFIERS.keys():
    for xy in ["_X", "_Y"]:
        column_names.append(id_name + xy)

for lr in ["_Right", "_Left"]:
    for id_name in HAND_IDENTIFIERS.keys():
        for xy in ["_X", "_Y"]:
            column_names.append(id_name + lr + xy)

column_names.append('labels')


def create_df(flnm, column_names):
    df = pd.DataFrame(columns=column_names)
    return df


def save_data(df, data, flnm):
    df = df.append(data.get_series(), ignore_index=True)
    df.to_pickle(flnm)


def obtain_pose_data(path):
    cap = cv2.VideoCapture(path)
    data = mp_holistic_data(column_names)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        # Recolor image to RGB
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Make detection
        holistic_results = holistic.process(image)
        # Extract feature and save to mp_pose_data class
        data.extract_data(holistic_results)
    cap.release()

    return data


if __name__ == '__main__':
    pass