forked from filby89/spectre
-
Notifications
You must be signed in to change notification settings - Fork 0
/
config.py
182 lines (142 loc) · 6.81 KB
/
config.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
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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
'''
Default config for SPECTRE - adapted from DECA
'''
from yacs.config import CfgNode as CN
import argparse
import yaml
import os
cfg = CN()
cfg.project_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), 'src', '..'))
cfg.device = 'cuda'
cfg.device_ids = '0'
cfg.pretrained_modelpath = os.path.join(cfg.project_dir, 'data', 'deca_model.tar')
cfg.output_dir = ''
cfg.rasterizer_type = 'pytorch3d'
# ---------------------------------------------------------------------------- #
# Options for FLAME and from original DECA
# ---------------------------------------------------------------------------- #
cfg.model = CN()
cfg.model.topology_path = os.path.join(cfg.project_dir, 'data' , 'head_template.obj')
# texture data original from http://files.is.tue.mpg.de/tbolkart/FLAME/FLAME_texture_data.zip
cfg.model.dense_template_path = os.path.join(cfg.project_dir, 'data', 'texture_data_256.npy')
cfg.model.fixed_displacement_path = os.path.join(cfg.project_dir, 'data', 'fixed_displacement_256.npy')
cfg.model.flame_model_path = os.path.join(cfg.project_dir, 'data', 'FLAME2020', 'generic_model.pkl')
cfg.model.flame_lmk_embedding_path = os.path.join(cfg.project_dir, 'data', 'landmark_embedding.npy')
cfg.model.face_mask_path = os.path.join(cfg.project_dir, 'data', 'uv_face_mask.png')
cfg.model.face_eye_mask_path = os.path.join(cfg.project_dir, 'data', 'uv_face_eye_mask.png')
cfg.model.mean_tex_path = os.path.join(cfg.project_dir, 'data', 'mean_texture.jpg')
cfg.model.tex_path = os.path.join(cfg.project_dir, 'data', 'FLAME_albedo_from_BFM.npz')
cfg.model.tex_type = 'BFM' # BFM, FLAME, albedoMM
cfg.model.uv_size = 256
cfg.model.param_list = ['shape', 'tex', 'exp', 'pose', 'cam', 'light']
cfg.model.n_shape = 100
cfg.model.n_tex = 50
cfg.model.n_exp = 50
cfg.model.n_cam = 3
cfg.model.n_pose = 6
cfg.model.n_light = 27
cfg.model.jaw_type = 'aa' # default use axis angle, another option: euler. Note that: aa is not stable in the beginning
cfg.model.model_type = "SPECTRE"
cfg.model.temporal = True
# ---------------------------------------------------------------------------- #
# Options for Dataset
# ---------------------------------------------------------------------------- #
cfg.dataset = CN()
cfg.dataset.LRS3_path = "/gpu-data3/filby/LRS3"
cfg.dataset.LRS3_landmarks_path = "../Visual_Speech_Recognition_for_Multiple_Languages/landmarks/LRS3/LRS3_landmarks"
cfg.dataset.LRS3_path = "/gpu-data3/filby/LRS3"
cfg.dataset.LRS3_landmarks_path = "../Visual_Speech_Recognition_for_Multiple_Languages/landmarks/LRS3/LRS3_landmarks"
cfg.dataset.LRS3_path = "/gpu-data3/filby/LRS3"
cfg.dataset.LRS3_landmarks_path = "../Visual_Speech_Recognition_for_Multiple_Languages/landmarks/LRS3/LRS3_landmarks"
cfg.dataset.batch_size = 1
cfg.dataset.K = 20
cfg.dataset.num_workers = 8
cfg.dataset.image_size = 224
cfg.dataset.scale_min = 1.4
cfg.dataset.scale_max = 1.8
cfg.dataset.trans_scale = 0.
cfg.dataset.fps = 25
cfg.dataset.test_datasets = ['LRS3']
# ---------------------------------------------------------------------------- #
# Options for training
# ---------------------------------------------------------------------------- #
cfg.train = CN()
cfg.train.max_epochs = 6
cfg.train.log_dir = 'logs'
cfg.train.log_steps = 10
cfg.train.vis_dir = 'train_images'
cfg.train.vis_steps = 500
cfg.train.write_summary = True
cfg.train.checkpoint_steps = 10000
cfg.train.val_vis_dir = 'val_images'
cfg.train.evaluation_steps = 10000
# ---------------------------------------------------------------------------- #
# Options for Losses
# ---------------------------------------------------------------------------- #
cfg.loss = CN()
cfg.loss.train = CN()
cfg.model.use_tex = True
cfg.model.regularization_type = 'nonlinear'
cfg.model.backbone = 'mobilenetv2' # perceptual encoder backbone
cfg.loss.train.landmark = 50
cfg.loss.train.lip_landmarks = 0
cfg.loss.train.relative_landmark = 50# 50
cfg.loss.train.photometric_texture = 0
cfg.loss.train.lipread = 2
cfg.loss.train.jaw_reg = 200
cfg.train.lr = 5e-5
cfg.loss.train.expression = 0.5
cfg.test_mode = False
def get_cfg_defaults():
"""Get a yacs CfgNode object with default values for my_project."""
# Return a clone so that the defaults will not be altered
# This is for the "local variable" use pattern
return cfg.clone()
def update_cfg(cfg, cfg_file):
cfg.merge_from_file(cfg_file)
return cfg.clone()
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--output_dir', type=str, help='output path')
parser.add_argument('--LRS3_path', default=None, type=str, help='path to LRS3 dataset')
parser.add_argument('--LRS3_landmarks_path', default=None, type=str, help='path to LRS3 landmarks')
parser.add_argument('--model_path', default=None, help='path to pretrained model')
parser.add_argument('--batch-size', type=int, default=1, help='the batch size')
parser.add_argument('--epochs', type=int, default=6, help='number of epochs to train for')
parser.add_argument('--K', type=int, default=20, help='length of sampled frame sequence')
parser.add_argument('--lipread', type=float, default=None, help='lipread loss weight')
parser.add_argument('--expression', type=float, default=None, help='expression loss weight')
parser.add_argument('--lr', type=float, default=None, help='learning rate')
parser.add_argument('--landmark', type=float, default=None, help='landmark loss weight')
parser.add_argument('--relative_landmark', type=float, default=None, help='relative landmark loss weight')
parser.add_argument('--backbone', type=str, default='mobilenetv2', choices=['mobilenetv2', 'resnet50'])
parser.add_argument('--test', action='store_true', help='test mode')
parser.add_argument('--test_datasets', type=str, nargs='+', default=['LRS3'], help='test datasets')
args = parser.parse_args()
cfg = get_cfg_defaults()
cfg.output_dir = args.output_dir
if args.model_path is not None:
cfg.pretrained_modelpath = args.model_path
if args.batch_size is not None:
cfg.dataset.batch_size = args.batch_size
cfg.dataset.K = args.K
if args.landmark is not None:
cfg.loss.train.landmark = args.landmark
if args.relative_landmark is not None:
cfg.loss.train.relative_landmark = args.relative_landmark
if args.lipread is not None:
cfg.loss.train.lipread = args.lipread
if args.expression is not None:
cfg.loss.train.expression = args.expression
if args.lr is not None:
cfg.train.lr = args.lr
if args.epochs is not None:
cfg.train.max_epochs = args.epochs
if args.LRS3_path is not None:
cfg.dataset.LRS3_path = args.LRS3_path
if args.LRS3_landmarks_path is not None:
cfg.dataset.LRS3_landmarks_path = args.LRS3_landmarks_path
cfg.model.backbone = args.backbone
cfg.test_mode = args.test
cfg.test_datasets = args.test_datasets
return cfg