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-06 09:55:48 +00:00
parent dfdf8d408e
commit 0b518e7f0e
46 changed files with 3392 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_curv.py
+17 -75
View File
@@ -251,90 +251,30 @@ def plot_everything(batch: dict, type: str):
def sample(folder):
load_model(model,folder)
def sample():
iteration = load_model(model, f'checkpoints/iccv/cvae_curv', return_iteration=True)
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples = 16
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)
ce_uv_loss = scce(tf.tile(batch['uv_curvature'][0][tf.newaxis,:,:], [nr_of_samples,1,1]), pred)
mesh_pred = uv2mesh(pred).numpy().argmax(axis=-1)
mesh_orig = batch['mesh_curvature'][0].numpy()
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
mesh_pred = [mesh_pred[i, mask] for i in range(mesh_pred.shape[0])]
curv_color = [np.zeros((mesh_pred[i].shape[0], 3)) for i in range(len(mesh_pred))]
mesh_orig = mesh_orig[mask] + 1
# 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)])
curv_color_gt = np.zeros((mesh_pred[0].shape[0], 3))
for i in range(6+1):
for j in range(len(mesh_pred)):
curv_color[j][mesh_pred[j] == i, :] = colors[i]
curv_color_gt[mesh_orig == i, :] = colors[i]
mesh_base = batch['mesh_posed'][0].numpy()[mask, :]
meshes = []
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=mesh_base[:,0],
y=mesh_base[:,1],
z=mesh_base[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
vertexcolor=curv_color_gt,
text='gt'
))
for i in range(nr_of_samples):
x = mesh_base.copy()
x[:, 1] += 1 * (i + 1)
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
vertexcolor=curv_color[i]
))
fig = go.Figure(data=meshes)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,nr_of_samples + 1],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
# fig.show()
pred = pred.numpy().argmax(-1)
frames = [Image.fromarray(pred[i]/6*255) for i in range(pred.shape[0])]
frames[0].save("temp/curv2.gif", format="GIF", append_images=frames[1:],
save_all=True, duration=500, loop=0)
color_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
for i in range(7):
color_img[img == i, :] = colors[i]
print()
if __name__ == '__main__':
init_wandb()
# init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
@@ -372,6 +312,8 @@ if __name__ == '__main__':
model = CVAE(config)
sample()
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=config['lr']) #0.0001
g_balancing_coeff = config['starting_balancing_coeff']