pose2sim/Pose2Sim/Utilities/Blazepose_runsave.py
2023-09-21 17:39:28 +02:00

243 lines
9.7 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
########################################################
## Run BlazePose and save coordinates ##
########################################################
Runs BlazePose (Mediapipe) on a video
Saves coordinates to OpenPose format (json files) or DeepLabCut format (csv or h5 table)
Optionally displays and saves images with keypoints overlayed
N.B.: First install mediapipe: `pip install mediapipe`
You may also need to install tables: `pip install tables`
Usage:
python -m Blazepose_runsave -i input_file --display --save_images --save_video --to_csv --to_h5 --to_json --model_complexity 2 -o output_folder
OR python -m Blazepose_runsave -i input_file --display --to_json --save_images
OR python -m Blazepose_runsave -i input_file -dJs
OR from Pose2Sim.Utilities import Blazepose_runsave; Blazepose_runsave.blazepose_detec_func(input_file=r'input_file', save_images=True, to_json=True, model_complexity=2)
'''
## INIT
import cv2
import mediapipe as mp
import os
import pandas as pd
import numpy as np
import json
import argparse
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
mp_pose = mp.solutions.pose
## AUTHORSHIP INFORMATION
__author__ = "David Pagnon"
__copyright__ = "Copyright 2023, Pose2Sim"
__credits__ = ["David Pagnon"]
__license__ = "BSD 3-Clause License"
__version__ = '0.4'
__maintainer__ = "David Pagnon"
__email__ = "contact@david-pagnon.com"
__status__ = "Development"
## FUNCTIONS
def save_to_csv_or_h5(kpt_list, output_folder, video_name, to_csv, to_h5):
'''
Saves blazepose keypoint coordinates to csv or h5 file,
in the DeepLabCut format.
INPUTS:
- kpt_list: List of lists of keypoints X and Y coordinates and likelihood, for each frame
- output_folder: Folder where to save the csv or h5 file
- video_name: Name of the video
- to_csv: Boolean, whether to save to csv
- to_h5: Boolean, whether to save to h5
OUTPUTS:
- Creation of csv or h5 file in output_folder
'''
# Prepare dataframe file
scorer = ['DavidPagnon']*len(mp_pose.PoseLandmark)*3
individuals = ['person']*len(mp_pose.PoseLandmark)*3
bodyparts = [[p.name]*3 for p in mp_pose.PoseLandmark]
bodyparts = [item for sublist in bodyparts for item in sublist]
coords = ['x', 'y', 'likelihood']*len(mp_pose.PoseLandmark)
tuples = list(zip(scorer, individuals, bodyparts, coords))
index_csv = pd.MultiIndex.from_tuples(tuples, names=['scorer', 'individuals', 'bodyparts', 'coords'])
df = pd.DataFrame(np.array(kpt_list).T, index=index_csv).T
if to_csv:
csv_file = os.path.join(output_folder, video_name+'.csv')
df.to_csv(csv_file, sep=',', index=True, line_terminator='\n')
if to_h5:
h5_file = os.path.join(output_folder, video_name+'.h5')
df.to_hdf(h5_file, index=True, key='blazepose_detection')
def save_to_json(kpt_list, output_folder, video_name):
'''
Saves blazepose keypoint coordinates to json file,
in the OpenPose format.
INPUTS:
- kpt_list: List of lists of keypoints X and Y coordinates and likelihood, for each frame
- output_folder: Folder where to save the csv or h5 file
- video_name: Name of the video
OUTPUTS:
- Creation of json files in output_folder/json_folder
'''
json_folder = os.path.join(output_folder, 'blaze_'+video_name + '_json')
if not os.path.exists(json_folder):
os.mkdir(json_folder)
print(json_folder)
# json preparation
json_dict = {'version':1.3, 'people':[]}
json_dict['people'] = [{'person_id':[-1],
'pose_keypoints_2d': [],
'face_keypoints_2d': [],
'hand_left_keypoints_2d':[],
'hand_right_keypoints_2d':[],
'pose_keypoints_3d':[],
'face_keypoints_3d':[],
'hand_left_keypoints_3d':[],
'hand_right_keypoints_3d':[]}]
# write each h5 line in json file
for frame, kpt in enumerate(kpt_list):
json_dict['people'][0]['pose_keypoints_2d'] = kpt
json_file = os.path.join(json_folder, 'blaze_'+video_name+'.'+str(frame).zfill(5)+'.json')
with open(json_file, 'w') as js_f:
js_f.write(json.dumps(json_dict))
def blazepose_detec_func(**args):
'''
Runs BlazePose (Mediapipe) on a video
Saves coordinates to OpenPose format (json files) or DeepLabCut format (csv or h5 table)
Optionally displays and saves images with keypoints overlayed
N.B.: First install mediapipe: `pip install mediapipe`
You may also need to install tables: `pip install tables`
Usage:
python -m Blazepose_runsave -i input_file --display --save_images --save_video --to_csv --to_h5 --to_json --model_complexity 2 -o output_folder
OR python -m Blazepose_runsave -i input_file --display --to_json --save_images
OR python -m Blazepose_runsave -i input_file -dJs
OR from Pose2Sim.Utilities import Blazepose_runsave; Blazepose_runsave.blazepose_detec_func(input_file=r'input_file', save_images=True, to_json=True, model_complexity=2)
'''
# Retrieve arguments
video_input = os.path.realpath(args.get('input_file'))
video_dir = os.path.dirname(video_input)
video_name = os.path.splitext(os.path.basename(video_input))[0]
output_folder = args.get('output_folder')
display = args.get('display')
save_images = args.get('save_images')
save_video = args.get('save_video')
to_csv = args.get('to_csv')
to_h5 = args.get('to_h5')
to_json = args.get('to_json')
model_complexity = int(args.get('model_complexity'))
if 'model_complexity' not in vars(): model_complexity=2
if to_csv or to_h5 or to_json or save_images or save_video:
if output_folder == None:
output_folder = video_dir
if not os.path.exists(os.path.realpath(output_folder)):
os.mkdir(os.path.realpath(output_folder))
# Run Blazepose
cap = cv2.VideoCapture(video_input)
W, H = cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
fps = cap.get(cv2.CAP_PROP_FPS)
count = 0
kpt_list = []
with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5, model_complexity=model_complexity) as pose:
while cap.isOpened():
ret, frame = cap.read()
if ret == True:
# Blazepose detection
results = pose.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
try:
kpt = [[p.x*W, p.y*H, p.visibility] for p in results.pose_landmarks.landmark]
kpt = [item for sublist in kpt for item in sublist]
mp_drawing.draw_landmarks(frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS, landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style())
except:
print(f'No person detected by BlazePose on frame {count}')
kpt=[np.nan]*3*33
# Display images
if display:
cv2.imshow('frame', frame)
if cv2.waitKey(30) & 0xFF == ord('q'):
break
# Save images
if save_images:
images_folder = os.path.join(output_folder, 'blaze_'+video_name + '_img')
if not os.path.exists(images_folder):
os.mkdir(images_folder)
cv2.imwrite(os.path.join(images_folder, 'blaze_'+video_name+'.'+str(count).zfill(5)+'.png'), frame)
# Save video
if save_video:
if count == 0:
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
writer = cv2.VideoWriter(os.path.join(output_folder, video_name+'_blaze.mp4'), fourcc, fps, (int(W), int(H)))
writer.write(frame)
# Store coordinates
if to_csv or to_h5 or to_json:
kpt_list.append(kpt)
count += 1
else:
break
cap.release()
if save_video:
writer.release()
cv2.destroyAllWindows()
# Save coordinates
if to_csv or to_h5:
save_to_csv_or_h5(kpt_list, output_folder, video_name, to_csv, to_h5)
if to_json:
save_to_json(kpt_list, output_folder, video_name)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_file', required = True, help='input video file')
parser.add_argument('-C', '--to_csv', required=False, action='store_true', help='save coordinates to csv')
parser.add_argument('-H', '--to_h5', required=False, action='store_true', help='save coordinates to h5')
parser.add_argument('-J', '--to_json', required=False, action='store_true', help='save coordinates to json')
parser.add_argument('-d', '--display', required = False, action='store_true', help='display images with overlayed coordinates')
parser.add_argument('-s', '--save_images', required = False, action='store_true', help='save images with overlayed coordinates')
parser.add_argument('-v', '--save_video', required = False, action='store_true', help='save video with overlayed coordinates')
parser.add_argument('-m', '--model_complexity', required = False, default = 2, help='model complexity. 0: fastest but less accurate, 2: most accurate but slowest')
parser.add_argument('-o', '--output_folder', required=False, help='output folder for coordinates and images')
args = vars(parser.parse_args())
blazepose_detec_func(**args)