luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

save

Browse Source
This commit is contained in:
luni
2023-07-14 10:12:00 +02:00
parent bb16aaf868
commit 88531da784
16 changed files with 411 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/get_final_cape.py
+18 -21
View File
@@ -13,10 +13,10 @@ parser.add_argument('-g', '--gpu', type=str, default='1')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_f
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
config_f = config_module.config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu']
@@ -54,13 +54,13 @@ from models.cvae_model import CVAE
from data.utils import Frame
uv2mesh = UV2Mesh(res=256, is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape'])
uv2mesh = UV2Mesh(res=256, is_skirt=config_f['is_skirt'], is_smpl_orig=config_f['is_cape'])
def step_pyramid(batch):
model.eval()
with torch.inference_mode():
levels_input, conds = group_inputs_for_levels(batch, config_1)
levels_input, conds = group_inputs_for_levels(batch, config_f)
out = []
@@ -70,16 +70,13 @@ def step_pyramid(batch):
mesh_pred = uv2mesh(pred)
if config_1['use_skinning']:
if config_f['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
if not config_1['is_cape']:
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
else:
mesh_pred_posed = mesh_pred_shaped
mesh_pred_posed = mesh_pred_shaped
else:
mesh_pred_posed = mesh_pred
@@ -94,8 +91,8 @@ cape_mask = np.zeros((6890), bool)
cape_mask[clothing_verts_idx] = True
def get_error(pred, frame: Frame):
gt = np.load(f'{config_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
pose_shape = np.load(f'{config_1["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
gt = np.load(f'{config_f["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape']
pose_shape = np.load(f'{config_f["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32)
pose, shape = pose_shape[:72], pose_shape[72:]
gender = int(shape[10] >= 0.5)
shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender].to('cpu'), dims=[[1], [2]])[0,smpl.to_keep,:].numpy()
@@ -155,30 +152,30 @@ def generate():
if __name__ == '__main__':
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
skirt_modifier = '_skirt' if config_1['is_skirt'] else ''
smpl = SMPL(orig_smpl=config_f['is_cape'], path='smpl', device=config_f['device'])
skirt_modifier = '_skirt' if config_f['is_skirt'] else ''
dataset_resolutions = []
if 'pyramid' in config_1:
dataset_resolutions = config_1['pyramid']['resolutions'].copy()
if 'pyramid' in config_f:
dataset_resolutions = config_f['pyramid']['resolutions'].copy()
if 512 not in dataset_resolutions:
dataset_resolutions.append(512)
# remove all values from config['signature] that start with mesh_
config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
config_f['signature'] = set([x for x in config_f['signature'] if not x.startswith('mesh_')])
dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
dataset_test = Dataset(config_f['test_dir'], config_f['signature'], device=config_f['device'], frames_file=config_f['test_frames'], use_cache=config_f['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_f['is_cape'])
# dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
is_pyramid = config_1['name'].startswith('pyr')
is_pyramid = config_f['name'].startswith('pyr')
model: Pyramid = Pyramid(config_1)
model: Pyramid = Pyramid(config_f)
it = load_model(model, 'checkpoints/iccvw/' + config_1['name'], iteration=args['iteration'])
it = load_model(model, 'checkpoints/iccvw/' + config_f['name'], iteration=args['iteration'])
model.eval()
# it = load_model(model, 'checkpoints/iccvw/' + config['name'])
name = config_1['name'] + '_' + str(it)
name = config_f['name'] + '_' + str(it)
folder = '/data-local/127_2/results'
out_folder = f'{folder}/{name}'