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

replace curvature with normals

Browse Source
This commit is contained in:
luni
2023-06-01 08:41:53 +00:00
parent a1470cfd6e
commit 534fbfcde1
4 changed files with 214 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cvae.py
+12 -10
View File
@@ -13,7 +13,8 @@ args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']:
pass
from configs.abl_norm import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
@@ -83,7 +84,7 @@ def train_step(batch):
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
model.train()
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
if config['use_skinning']:
@@ -139,7 +140,7 @@ def valid_step(batch):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
if config['use_skinning']:
@@ -188,15 +189,16 @@ def plot_everything(batch: dict, type: str):
model.eval()
with torch.inference_mode():
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
mesh_pred = uv2mesh(pred).detach().cpu().numpy()
orig_img = torch.permute(uv_orig, [0,2,3,1])
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
pred_img = torch.permute(pred, [0,2,3,1])
img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
img = torch.concat((uv_static_img, uv_body_posed_img, uv_normals_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -262,10 +264,10 @@ def sample(folder):
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_normals = tf.one_hot(batch['uv_normals'], 7)
nr_of_samples=2
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
@@ -281,7 +283,7 @@ def sample(folder):
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_normals[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
@@ -457,7 +459,7 @@ def train():
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
ckpt_saver(model, step, optimizer)
ckpt_saver(model, step, optimizer, None)
step += 1
@@ -491,6 +493,6 @@ if __name__ == '__main__':
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
train()