sapiens-demo / config.py
fiesty-bear's picture
added skeleton
47b2ed0 verified
raw
history blame
12.4 kB
SAPIENS_LITE_MODELS_URL = {
"depth": {
"sapiens_0.3b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_0.3b/sapiens_0.3b_render_people_epoch_100_torchscript.pt2?download=true",
"sapiens_0.6b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_0.6b/sapiens_0.6b_render_people_epoch_70_torchscript.pt2?download=true",
"sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_1b/sapiens_1b_render_people_epoch_88_torchscript.pt2?download=true",
"sapiens_2b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_2b/sapiens_2b_render_people_epoch_25_torchscript.pt2?download=true"
},
"detector": {},
"normal": {
"sapiens_0.3b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_0.3b/sapiens_0.3b_normal_render_people_epoch_66_torchscript.pt2?download=true",
"sapiens_0.6b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_0.6b/sapiens_0.6b_normal_render_people_epoch_200_torchscript.pt2?download=true",
"sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_1b/sapiens_1b_normal_render_people_epoch_115_torchscript.pt2?download=true",
"sapiens_2b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_2b/sapiens_2b_normal_render_people_epoch_70_torchscript.pt2?download=true"
},
"pose": {
"sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/pose/checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_AP_640_torchscript.pt2?download=true"
},
"seg": {
"sapiens_0.3b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.3b/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2?download=true",
"sapiens_0.6b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.6b/sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2?download=true",
"sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2?download=true",
"sapiens_2b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_2b/sapiens_2b_goliath_best_goliath_mIoU_8179_epoch_181_torchscript.pt2?download=true"
}
}
SAPIENS_LITE_MODELS_PATH = {
"depth": {
"sapiens_0.3b": "checkpoints/depth/sapiens_0.3b_torchscript.pt2",
"sapiens_0.6b": "checkpoints/depth/sapiens_0.6b_torchscript.pt2",
"sapiens_1b": "checkpoints/depth/sapiens_1b_torchscript.pt2",
"sapiens_2b": "checkpoints/depth/sapiens_2b_torchscript.pt2"
},
"detector": {},
"normal": {
"sapiens_0.3b": "checkpoints/normal/sapiens_0.3b_torchscript.pt2",
"sapiens_0.6b": "checkpoints/normal/sapiens_0.6b_torchscript.pt2",
"sapiens_1b": "checkpoints/normal/sapiens_1b_torchscript.pt2",
"sapiens_2b": "checkpoints/normal/sapiens_2b_torchscript.pt2"
},
"pose": {
"sapiens_1b": "checkpoints/pose/sapiens_1b_torchscript.pt2"
},
"seg": {
"sapiens_0.3b": "checkpoints/seg/sapiens_0.3b_torchscript.pt2",
"sapiens_0.6b": "checkpoints/seg/sapiens_0.6b_torchscript.pt2",
"sapiens_1b": "checkpoints/seg/sapiens_1b_torchscript.pt2",
"sapiens_2b": "checkpoints/seg/sapiens_2b_torchscript.pt2"
}
}
LABELS_TO_IDS = {
"Background": 0,
"Apparel": 1,
"Face Neck": 2,
"Hair": 3,
"Left Foot": 4,
"Left Hand": 5,
"Left Lower Arm": 6,
"Left Lower Leg": 7,
"Left Shoe": 8,
"Left Sock": 9,
"Left Upper Arm": 10,
"Left Upper Leg": 11,
"Lower Clothing": 12,
"Right Foot": 13,
"Right Hand": 14,
"Right Lower Arm": 15,
"Right Lower Leg": 16,
"Right Shoe": 17,
"Right Sock": 18,
"Right Upper Arm": 19,
"Right Upper Leg": 20,
"Torso": 21,
"Upper Clothing": 22,
"Lower Lip": 23,
"Upper Lip": 24,
"Lower Teeth": 25,
"Upper Teeth": 26,
"Tongue": 27,
}
skeleton_info={
0:
dict(link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
1:
dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255, 0]),
2:
dict(link=('right_ankle', 'right_knee'), id=2, color=[255, 128, 0]),
3:
dict(link=('right_knee', 'right_hip'), id=3, color=[255, 128, 0]),
4:
dict(link=('left_hip', 'right_hip'), id=4, color=[51, 153, 255]),
5:
dict(link=('left_shoulder', 'left_hip'), id=5, color=[51, 153, 255]),
6:
dict(link=('right_shoulder', 'right_hip'), id=6, color=[51, 153, 255]),
7:
dict(
link=('left_shoulder', 'right_shoulder'),
id=7,
color=[51, 153, 255]),
8:
dict(link=('left_shoulder', 'left_elbow'), id=8, color=[0, 255, 0]),
9:
dict(
link=('right_shoulder', 'right_elbow'), id=9, color=[255, 128, 0]),
10:
dict(link=('left_elbow', 'left_wrist'), id=10, color=[0, 255, 0]),
11:
dict(link=('right_elbow', 'right_wrist'), id=11, color=[255, 128, 0]),
12:
dict(link=('left_eye', 'right_eye'), id=12, color=[51, 153, 255]),
13:
dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
14:
dict(link=('nose', 'right_eye'), id=14, color=[51, 153, 255]),
15:
dict(link=('left_eye', 'left_ear'), id=15, color=[51, 153, 255]),
16:
dict(link=('right_eye', 'right_ear'), id=16, color=[51, 153, 255]),
17:
dict(link=('left_ear', 'left_shoulder'), id=17, color=[51, 153, 255]),
18:
dict(
link=('right_ear', 'right_shoulder'), id=18, color=[51, 153, 255]),
19:
dict(link=('left_ankle', 'left_big_toe'), id=19, color=[0, 255, 0]),
20:
dict(link=('left_ankle', 'left_small_toe'), id=20, color=[0, 255, 0]),
21:
dict(link=('left_ankle', 'left_heel'), id=21, color=[0, 255, 0]),
22:
dict(
link=('right_ankle', 'right_big_toe'), id=22, color=[255, 128, 0]),
23:
dict(
link=('right_ankle', 'right_small_toe'),
id=23,
color=[255, 128, 0]),
24:
dict(link=('right_ankle', 'right_heel'), id=24, color=[255, 128, 0]),
25:
dict(
link=('left_hand_root', 'left_thumb1'), id=25, color=[255, 128,
0]),
26:
dict(link=('left_thumb1', 'left_thumb2'), id=26, color=[255, 128, 0]),
27:
dict(link=('left_thumb2', 'left_thumb3'), id=27, color=[255, 128, 0]),
28:
dict(link=('left_thumb3', 'left_thumb4'), id=28, color=[255, 128, 0]),
29:
dict(
link=('left_hand_root', 'left_forefinger1'),
id=29,
color=[255, 153, 255]),
30:
dict(
link=('left_forefinger1', 'left_forefinger2'),
id=30,
color=[255, 153, 255]),
31:
dict(
link=('left_forefinger2', 'left_forefinger3'),
id=31,
color=[255, 153, 255]),
32:
dict(
link=('left_forefinger3', 'left_forefinger4'),
id=32,
color=[255, 153, 255]),
33:
dict(
link=('left_hand_root', 'left_middle_finger1'),
id=33,
color=[102, 178, 255]),
34:
dict(
link=('left_middle_finger1', 'left_middle_finger2'),
id=34,
color=[102, 178, 255]),
35:
dict(
link=('left_middle_finger2', 'left_middle_finger3'),
id=35,
color=[102, 178, 255]),
36:
dict(
link=('left_middle_finger3', 'left_middle_finger4'),
id=36,
color=[102, 178, 255]),
37:
dict(
link=('left_hand_root', 'left_ring_finger1'),
id=37,
color=[255, 51, 51]),
38:
dict(
link=('left_ring_finger1', 'left_ring_finger2'),
id=38,
color=[255, 51, 51]),
39:
dict(
link=('left_ring_finger2', 'left_ring_finger3'),
id=39,
color=[255, 51, 51]),
40:
dict(
link=('left_ring_finger3', 'left_ring_finger4'),
id=40,
color=[255, 51, 51]),
41:
dict(
link=('left_hand_root', 'left_pinky_finger1'),
id=41,
color=[0, 255, 0]),
42:
dict(
link=('left_pinky_finger1', 'left_pinky_finger2'),
id=42,
color=[0, 255, 0]),
43:
dict(
link=('left_pinky_finger2', 'left_pinky_finger3'),
id=43,
color=[0, 255, 0]),
44:
dict(
link=('left_pinky_finger3', 'left_pinky_finger4'),
id=44,
color=[0, 255, 0]),
45:
dict(
link=('right_hand_root', 'right_thumb1'),
id=45,
color=[255, 128, 0]),
46:
dict(
link=('right_thumb1', 'right_thumb2'), id=46, color=[255, 128, 0]),
47:
dict(
link=('right_thumb2', 'right_thumb3'), id=47, color=[255, 128, 0]),
48:
dict(
link=('right_thumb3', 'right_thumb4'), id=48, color=[255, 128, 0]),
49:
dict(
link=('right_hand_root', 'right_forefinger1'),
id=49,
color=[255, 153, 255]),
50:
dict(
link=('right_forefinger1', 'right_forefinger2'),
id=50,
color=[255, 153, 255]),
51:
dict(
link=('right_forefinger2', 'right_forefinger3'),
id=51,
color=[255, 153, 255]),
52:
dict(
link=('right_forefinger3', 'right_forefinger4'),
id=52,
color=[255, 153, 255]),
53:
dict(
link=('right_hand_root', 'right_middle_finger1'),
id=53,
color=[102, 178, 255]),
54:
dict(
link=('right_middle_finger1', 'right_middle_finger2'),
id=54,
color=[102, 178, 255]),
55:
dict(
link=('right_middle_finger2', 'right_middle_finger3'),
id=55,
color=[102, 178, 255]),
56:
dict(
link=('right_middle_finger3', 'right_middle_finger4'),
id=56,
color=[102, 178, 255]),
57:
dict(
link=('right_hand_root', 'right_ring_finger1'),
id=57,
color=[255, 51, 51]),
58:
dict(
link=('right_ring_finger1', 'right_ring_finger2'),
id=58,
color=[255, 51, 51]),
59:
dict(
link=('right_ring_finger2', 'right_ring_finger3'),
id=59,
color=[255, 51, 51]),
60:
dict(
link=('right_ring_finger3', 'right_ring_finger4'),
id=60,
color=[255, 51, 51]),
61:
dict(
link=('right_hand_root', 'right_pinky_finger1'),
id=61,
color=[0, 255, 0]),
62:
dict(
link=('right_pinky_finger1', 'right_pinky_finger2'),
id=62,
color=[0, 255, 0]),
63:
dict(
link=('right_pinky_finger2', 'right_pinky_finger3'),
id=63,
color=[0, 255, 0]),
64:
dict(
link=('right_pinky_finger3', 'right_pinky_finger4'),
id=64,
color=[0, 255, 0])
},