-
Notifications
You must be signed in to change notification settings - Fork 1
/
demo.py
58 lines (41 loc) · 1.69 KB
/
demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
'''
@FileName : demo.py
@EditTime : 2024-01-15 14:17:21
@Author : Buzhen Huang
@Email : buzhenhuang@outlook.com
@Description :
'''
import torch
from cmd_parser import parse_config
from modules import init, DatasetLoader, ModelLoader
from yolox.yolox import Predictor
from alphapose_module.alphapose_core import AlphaPose_Predictor
###########Load config file in debug mode#########
# import sys
# sys.argv = ['','--config=cfg_files/demo.yaml']
def main(**args):
# Global setting
dtype = torch.float32
device = torch.device(index=args.get('gpu_index'), type='cuda')
# Initialize project setting, e.g., create output folder, load SMPL model
out_dir, logger, smpl, generator, occlusions = init(dtype=dtype, **args)
# human detection
yolox_model_dir = R'data/bytetrack_x_mot17.pth.tar'
yolox_thres = 0.3
yolox_predictor = Predictor(yolox_model_dir, yolox_thres)
# 2D pose estimation
alpha_config = R'alphapose_module/configs/halpe_26/resnet/256x192_res50_lr1e-3_1x.yaml'
alpha_checkpoint = R'data/halpe26_fast_res50_256x192.pth'
alpha_thres = 0.1
alpha_predictor = AlphaPose_Predictor(alpha_config, alpha_checkpoint, alpha_thres)
# load Pose2UV model
model = ModelLoader(device=device, out_dir=out_dir, smpl=smpl, generator=generator, **args)
# load data
dataset = DatasetLoader(smpl_model=smpl, generator=generator, occlusions=occlusions, **args)
test_dataset = dataset.load_demo_data()
task = args.pop('task')
exec('from process import %s' %task)
eval('%s' %task)(model, yolox_predictor, alpha_predictor, test_dataset, device=device)
if __name__ == "__main__":
args = parse_config()
main(**args)