93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
import os
|
|
import os.path as osp
|
|
import sys
|
|
import datetime
|
|
from mmengine.config import Config as MMConfig
|
|
|
|
import os
|
|
import os.path as osp
|
|
|
|
# write your path here
|
|
|
|
encoder_config_file = '/home/weichen/sst/smplerx/main/transformer_utils/configs/smpler_x/encoder/body_encoder_huge.py'
|
|
human_model_path = '/home/weichen/wc_workspace/models/human_model'
|
|
|
|
# <configure your smplerx model path here>
|
|
model_path_dict = {
|
|
'smpler_x_h32': '/home/weichen/wc_workspace/models/smplerx/smpler_x_h32_correct.pth.tar',
|
|
'smpler_x_l32': '',
|
|
'smpler_x_b32': '',
|
|
'smpler_x_s32': '',
|
|
}
|
|
|
|
|
|
class Config:
|
|
def get_config_fromfile(self, config_path):
|
|
|
|
self.config_path = config_path
|
|
cfg = MMConfig.fromfile(self.config_path)
|
|
|
|
# update config
|
|
cfg.encoder_config_file = encoder_config_file
|
|
cfg.human_model_path = human_model_path
|
|
|
|
self.__dict__.update(dict(cfg))
|
|
|
|
# update dir
|
|
self.cur_dir = osp.dirname(os.path.abspath(__file__))
|
|
self.root_dir = osp.join(self.cur_dir, '..')
|
|
self.data_dir = osp.join(self.root_dir, 'dataset')
|
|
self.human_model_path = human_model_path
|
|
|
|
## add some paths to the system root dir
|
|
sys.path.insert(0, osp.join(self.root_dir, 'common'))
|
|
|
|
|
|
def prepare_dirs(self, exp_name):
|
|
time_str = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
|
|
self.output_dir = osp.join(self.root_dir, f'{exp_name}_{time_str}')
|
|
self.model_dir = osp.join(self.output_dir, 'model_dump')
|
|
self.vis_dir = osp.join(self.output_dir, 'vis')
|
|
self.log_dir = osp.join(self.output_dir, 'log')
|
|
self.code_dir = osp.join(self.output_dir, 'code')
|
|
self.result_dir = osp.join(self.output_dir, 'result')
|
|
|
|
from utils.dir import make_folder
|
|
make_folder(self.model_dir)
|
|
make_folder(self.vis_dir)
|
|
make_folder(self.log_dir)
|
|
make_folder(self.code_dir)
|
|
make_folder(self.result_dir)
|
|
|
|
## copy some code to log dir as a backup
|
|
copy_files = ['main/train.py', 'main/test.py', 'common/base.py',
|
|
'common/nets', 'main/SMPLer_X.py',
|
|
'data/dataset.py', 'data/MSCOCO/MSCOCO.py', 'data/AGORA/AGORA.py']
|
|
for file in copy_files:
|
|
os.system(f'cp -r {self.root_dir}/{file} {self.code_dir}')
|
|
|
|
def update_test_config(self, testset, agora_benchmark, shapy_eval_split, pretrained_model_path, use_cache,
|
|
eval_on_train=False, vis=False):
|
|
self.testset = testset
|
|
self.agora_benchmark = agora_benchmark
|
|
self.pretrained_model_path = pretrained_model_path
|
|
self.shapy_eval_split = shapy_eval_split
|
|
self.use_cache = use_cache
|
|
self.eval_on_train = eval_on_train
|
|
self.vis = vis
|
|
|
|
|
|
def update_config(self, num_gpus, pretrained_model_path, output_folder, device):
|
|
self.num_gpus = num_gpus
|
|
self.pretrained_model_path = pretrained_model_path
|
|
self.log_dir = output_folder
|
|
self.device = device
|
|
|
|
# Save
|
|
cfg_save = MMConfig(self.__dict__)
|
|
|
|
cfg = Config()
|
|
cfg.human_model_path = human_model_path
|
|
|
|
|