HyperPose/python_demo.py at master · tensorlayer/HyperPose · GitHub

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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
#!/usr/bin/env python3
importos
importcv2
importsys
importglob
importargparse
importmatplotlib
importnumpyasnp
importtensorflowastf
importtensorlayerastl
fromhyperposeimportConfig,Model,Dataset

if__name__=='__main__':
parser=argparse.ArgumentParser(description='Hyperpose')
parser.add_argument("--model_type",
type=str,
default="Openpose",
help="human pose estimation model type, available options: Openpose, LightweightOpenpose ,MobilenetThinOpenpose, PoseProposal, Pifpaf")
parser.add_argument("--model_backbone",
type=str,
default="Default",
help="model backbone, available options: Mobilenet, Vggtiny, Vgg19, Resnet18, Resnet50")
parser.add_argument("--dataset_type",
type=str,
default="MSCOCO",
help="dataset name,to determine which dataset to use, available options: MSCOCO, MPII ")
parser.add_argument("--model_name",
type=str,
default="default_name",
help="model name,to distinguish model and determine model dir")
parser.add_argument("--image_dir",
type=str,
default="./save_dir/example_dir/image",
help="image paths to be processed by the model"
 )
parser.add_argument("--output_dir",
type=str,
default="./save_dir/example_dir/output_dir",
help="ouput directory of the model forwarding"
 )

args=parser.parse_args()
# config model
Config.set_model_name(args.model_name)
Config.set_model_type(Config.MODEL[args.model_type])
Config.set_dataset_type(Config.DATA[args.dataset_type])
Config.set_model_backbone(Config.BACKBONE[args.model_backbone])
config=Config.get_config()
output_dir=os.path.join(args.output_dir,args.model_name)
os.makedirs(output_dir, exist_ok=True)

# contruct model and processors
model=Model.get_model(config)
# visualizer
VisualizerClass=Model.get_visualizer(config)
visualizer=VisualizerClass(save_dir=output_dir, parts=model.parts, limbs=model.limbs)
# post processor
PostProcessorClass=Model.get_postprocessor(config)
post_processor=PostProcessorClass(parts=model.parts, limbs=model.limbs, hin=model.hin, win=model.win, hout=model.hout,
wout=model.wout, colors=model.colors)
# image processor
ImageProcessorClass=Model.get_imageprocessor()
image_processor=ImageProcessorClass(input_h=model.hin, input_w=model.win)

# load weights
model_weight_path=f"./save_dir/{args.model_name}/model_dir/newest_model.npz"
model.load_weights(model_weight_path, format="npz_dict")
model.eval()
# begin process
forimage_pathinglob.glob(f"{args.image_dir}/*"):
image_name=os.path.basename(image_path)
print(f"processing image:{image_name}")
# image read, normalize, and scale
image=image_processor.read_image_rgb_float(image_path)
input_image, scale, pad=image_processor.image_pad_and_scale(image)
input_image=np.transpose(input_image,[2,0,1])[np.newaxis,:,:,:]
# model forward
predict_x=model.forward(input_image)
# post process
humans=post_processor.process(predict_x)[0]
# visualize heatmaps
visualizer.visualize(image_batch=input_image, predict_x=predict_x, humans_list=[humans], name=image_name)
# visualize results (restore detected humans)
print(f"{len(humans)} humans detected")
forhuman_idx,humaninenumerate(humans,start=1):
human.unpad(pad)
human.unscale(scale)
print(f"human:{human_idx} num of detected body joints:{human.get_partnum()}")
human.print()
visualizer.visualize_result(image=image, humans=humans, name=f"{image_name}_result")