EasyMocap/apps/preprocess/extract_keypoints.py

'''
  @ 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()