''' Default config for SPECTRE - adapted from DECA ''' from yacs.config import CfgNode as CN import argparse import yaml import os from motiondiff_modules import download_models, CKPT_DIR_PATH 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(str(CKPT_DIR_PATH), '..' , 'smpl_models' , 'FLAME_NEUTRAL.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