197 lines
7.5 KiB
Python
197 lines
7.5 KiB
Python
'''
|
|
@ Date: 2021-08-19 22:06:22
|
|
@ Author: Qing Shuai
|
|
@ LastEditors: Qing Shuai
|
|
@ LastEditTime: 2022-05-23 22:58:57
|
|
@ FilePath: /EasyMocapPublic/apps/preprocess/extract_keypoints.py
|
|
'''
|
|
import os
|
|
from os.path import join
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
def load_subs(path, subs):
|
|
if len(subs) == 0:
|
|
subs = sorted(os.listdir(join(path, 'images')))
|
|
subs = [sub for sub in subs if os.path.isdir(join(path, 'images', sub))]
|
|
if len(subs) == 0:
|
|
subs = ['']
|
|
return subs
|
|
|
|
config = {
|
|
'openpose':{
|
|
'root': '',
|
|
'res': 1,
|
|
'hand': False,
|
|
'face': False,
|
|
'vis': False,
|
|
'ext': '.jpg'
|
|
},
|
|
'openposecrop': {},
|
|
'feet':{
|
|
'root': '',
|
|
'res': 1,
|
|
'hand': False,
|
|
'face': False,
|
|
'vis': False,
|
|
'ext': '.jpg'
|
|
},
|
|
'feetcrop':{
|
|
'root': '',
|
|
'res': 1,
|
|
'hand': False,
|
|
'face': False,
|
|
'vis': False,
|
|
'ext': '.jpg'
|
|
},
|
|
'yolo':{
|
|
'ckpt_path': 'data/models/yolov4.weights',
|
|
'conf_thres': 0.3,
|
|
'box_nms_thres': 0.5, # means keeping the bboxes that IOU<0.5
|
|
'ext': '.jpg',
|
|
'isWild': False,
|
|
},
|
|
'hrnet':{
|
|
'nof_joints': 17,
|
|
'c': 48,
|
|
'checkpoint_path': 'data/models/pose_hrnet_w48_384x288.pth'
|
|
},
|
|
'yolo-hrnet':{},
|
|
'mp-pose':{
|
|
'model_complexity': 2,
|
|
'min_detection_confidence':0.5,
|
|
'min_tracking_confidence': 0.5
|
|
},
|
|
'mp-holistic':{
|
|
'model_complexity': 2,
|
|
# 'refine_face_landmarks': True,
|
|
'min_detection_confidence':0.5,
|
|
'min_tracking_confidence': 0.5
|
|
},
|
|
'mp-handl':{
|
|
'model_complexity': 1,
|
|
'min_detection_confidence':0.3,
|
|
'min_tracking_confidence': 0.1,
|
|
'static_image_mode': False,
|
|
},
|
|
'mp-handr':{
|
|
'model_complexity': 1,
|
|
'min_detection_confidence':0.3,
|
|
'min_tracking_confidence': 0.1,
|
|
'static_image_mode': False,
|
|
},
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('path', type=str, default=None, help="the path of data")
|
|
parser.add_argument('--subs', type=str, nargs='+', default=[], help="the path of data")
|
|
# Output Control
|
|
parser.add_argument('--annot', type=str, default='annots',
|
|
help="sub directory name to store the generated annotation files, default to be annots")
|
|
# Detection Control
|
|
parser.add_argument('--mode', type=str, default='openpose', choices=[
|
|
'openpose', 'feet', 'feetcrop', 'openposecrop',
|
|
'yolo-hrnet', 'yolo', 'hrnet',
|
|
'mp-pose', 'mp-holistic', 'mp-handl', 'mp-handr', 'mp-face'],
|
|
help="model to extract joints from image")
|
|
# Openpose
|
|
parser.add_argument('--openpose', type=str,
|
|
default=os.environ.get('openpose', 'openpose'))
|
|
parser.add_argument('--openpose_res', type=float, default=1)
|
|
parser.add_argument('--gpus', type=int, default=[], nargs='+')
|
|
parser.add_argument('--ext', type=str, default='.jpg')
|
|
parser.add_argument('--tmp', type=str, default=os.path.abspath('tmp'))
|
|
parser.add_argument('--hand', action='store_true')
|
|
parser.add_argument('--face', action='store_true')
|
|
parser.add_argument('--wild', action='store_true',
|
|
help='remove crowd class of yolo')
|
|
parser.add_argument('--reverse', action='store_true')
|
|
parser.add_argument('--force', action='store_true')
|
|
args = parser.parse_args()
|
|
config['yolo']['isWild'] = args.wild
|
|
mode = args.mode
|
|
if not os.path.exists(join(args.path, 'images')) and os.path.exists(join(args.path, 'videos')):
|
|
# default extract image
|
|
cmd = f'''python3 apps/preprocess/extract_image.py {args.path}'''
|
|
os.system(cmd)
|
|
subs = load_subs(args.path, args.subs)
|
|
if len(args.gpus) != 0:
|
|
# perform multiprocess by runcmd
|
|
from easymocap.mytools.debug_utils import run_cmd
|
|
nproc = len(args.gpus)
|
|
plist = []
|
|
for i in range(nproc):
|
|
if len(subs[i::nproc]) == 0:
|
|
continue
|
|
cmd = f'export CUDA_VISIBLE_DEVICES={args.gpus[i]} && python3 apps/preprocess/extract_keypoints.py {args.path} --mode {args.mode} --subs {" ".join(subs[i::nproc])}'
|
|
cmd += f' --annot {args.annot} --ext {args.ext}'
|
|
if args.hand:
|
|
cmd += ' --hand'
|
|
if args.face:
|
|
cmd += ' --face'
|
|
if args.force:
|
|
cmd += ' --force'
|
|
cmd += f' --tmp {args.tmp}'
|
|
cmd += ' &'
|
|
print(cmd)
|
|
p = run_cmd(cmd, bg=False)
|
|
plist.extend(p)
|
|
for p in plist:
|
|
p.join()
|
|
exit()
|
|
global_tasks = []
|
|
if len(subs) == 0:
|
|
subs = ['']
|
|
for sub in subs:
|
|
config[mode]['force'] = args.force
|
|
image_root = join(args.path, 'images', sub)
|
|
annot_root = join(args.path, args.annot, sub)
|
|
tmp_root = join(args.path, mode, sub)
|
|
if os.path.exists(annot_root) and not args.force:
|
|
# check the number of annots and images
|
|
if len(os.listdir(image_root)) == len(os.listdir(annot_root)) and mode != 'feetcrop':
|
|
print('[Skip] detection {}'.format(annot_root))
|
|
continue
|
|
if mode == 'openpose':
|
|
from easymocap.estimator.openpose_wrapper import extract_2d
|
|
config[mode]['root'] = args.openpose
|
|
config[mode]['hand'] = args.hand
|
|
config[mode]['face'] = args.face
|
|
config[mode]['res'] = args.openpose_res
|
|
config[mode]['ext'] = args.ext
|
|
global_tasks = extract_2d(image_root, annot_root, tmp_root, config[mode])
|
|
elif mode == 'feet':
|
|
from easymocap.estimator.openpose_wrapper import FeetEstimator
|
|
config[mode]['openpose'] = args.openpose
|
|
estimator = FeetEstimator(openpose=args.openpose)
|
|
estimator.detect_foot(image_root, annot_root, args.ext)
|
|
elif mode == 'yolo' or mode == 'openposecrop':
|
|
from easymocap.estimator.yolohrnet_wrapper import extract_bbox
|
|
config['yolo']['force'] = False # donot redetect
|
|
config['yolo']['ext'] = args.ext
|
|
extract_bbox(image_root, annot_root, **config['yolo'])
|
|
elif mode == 'hrnet':
|
|
from easymocap.estimator.yolohrnet_wrapper import extract_hrnet
|
|
config[mode]['ext'] = args.ext
|
|
extract_hrnet(image_root, annot_root, **config[mode])
|
|
elif mode == 'yolo-hrnet':
|
|
from easymocap.estimator.yolohrnet_wrapper import extract_yolo_hrnet
|
|
extract_yolo_hrnet(image_root, annot_root, args.ext, config['yolo'], config['hrnet'])
|
|
elif mode in ['mp-pose', 'mp-holistic', 'mp-handl', 'mp-handr', 'mp-face']:
|
|
from easymocap.estimator.mediapipe_wrapper import extract_2d
|
|
config[mode]['ext'] = args.ext
|
|
extract_2d(image_root, annot_root, config[mode], mode=mode.replace('mp-', ''))
|
|
if mode == 'feetcrop' or mode == 'openposecrop':
|
|
from easymocap.estimator.openpose_wrapper import FeetEstimatorByCrop
|
|
config[mode]['openpose'] = args.openpose
|
|
estimator = FeetEstimatorByCrop(openpose=args.openpose,
|
|
tmpdir=join(args.tmp, '{}_{}'.format(os.path.basename(args.path), sub)),
|
|
fullbody=mode=='openposecrop',
|
|
hand=(mode=='openposecrop')or args.hand,
|
|
face=args.face)
|
|
estimator.detect_foot(image_root, annot_root, args.ext)
|
|
|
|
for task in global_tasks:
|
|
task.join() |