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

save

Browse Source
This commit is contained in:
luni
2023-03-07 08:59:20 +00:00
parent 0b518e7f0e
commit e192f6f9e3
21 changed files with 1004 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_curv.py
+42 -18
View File
@@ -153,11 +153,11 @@ def dice_dense(prediction, ground_truth, weight_map=None):
return tf.reduce_mean(dice_score)
# @tf.function
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
dice = dice_dense(pred, batch['uv_curvature'])
@@ -189,7 +189,7 @@ def train_step(batch):
@tf.function
def valid_step(batch):
pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
loss = ce_uv_loss
@@ -216,7 +216,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
pred = np.argmax(pred.numpy(), axis=-1)
mesh_pred = uv2mesh(pred)
@@ -252,27 +252,51 @@ def plot_everything(batch: dict, type: str):
def sample():
iteration = load_model(model, f'checkpoints/iccv/cvae_curv', return_iteration=True)
iteration = load_model(model, f'checkpoints/iccv/cvae_curv_3', return_iteration=True)
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
nr_of_samples = 4
# pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask'])
# pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pred = model.sample(tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1]))
# pred,_,_ = model([tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1])], training=False)
orig_img = np.argmax(batch['uv_curvature'][0].numpy(), axis=-1)
pred_img = np.argmax(pred.numpy(), axis=-1)
img = np.concatenate((orig_img, np.concatenate(pred_img, axis=1)), axis=1)
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
color_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
for i in range(7):
color_img[img == i, :] = colors[i]
for i in range(4):
# pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask'])
# pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
pred = model.sample(tf.tile(batch['uv_static'][i][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][i][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][i][None,:,:], [nr_of_samples,1,1]))
# pred,_,_ = model([tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1])], training=False)
# pred = model.call_sample([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
orig_img = np.argmax(batch['uv_curvature'][i].numpy(), axis=-1)
pred_img = np.argmax(pred.numpy(), axis=-1)
img = np.concatenate((orig_img, np.concatenate(pred_img, axis=1)), axis=1)
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
color_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
for i in range(7):
color_img[img == i, :] = colors[i]
ex.imshow(color_img).show()
print()
print()
def sample_2(batch):
nr_of_samples = 4
# pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask'])
# pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
# pred = model.sample(tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1]))
pred,_,_ = model([tf.tile(batch['uv_static'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_body_posed'][0][None,:,:,:], [nr_of_samples,1,1,1]),tf.tile(batch['uv_curvature'][0][None,:,:,:], [nr_of_samples,1,1,1]), tf.tile(batch['uv_mask'][0][None,:,:], [nr_of_samples,1,1])], training=False)
# pred = model.call_sample([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
orig_img = np.argmax(batch['uv_curvature'][0].numpy(), axis=-1)
pred_img = np.argmax(pred.numpy(), axis=-1)
img = np.concatenate((orig_img, np.concatenate(pred_img, axis=1)), axis=1)
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
color_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
for i in range(7):
color_img[img == i, :] = colors[i]
ex.imshow(color_img).show()
print()
if __name__ == '__main__':
# init_wandb()
parser = argparse.ArgumentParser(description='')