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

working weight learning

Browse Source
This commit is contained in:
luni
2023-02-15 11:19:14 +00:00
parent 8fbd075f92
commit 61a609c8f7
7 changed files with 96 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aaa_helpers.txt
+3
View File
@@ -1,3 +1,6 @@
nohup python cvae.py --name="lhunor_weights_test_2" > nohup/iccv/weights_test.out 2>&1 &
one final convolutional layer for smoothing, on the final coordinates
skirts!!!!
configs/config_abl_all_post.py
+5 -5
View File
@@ -7,7 +7,7 @@ config = dict()
# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
# DON'T FORGET TO CHECK AVG GARMENT USED
config['use_skinning'] = True # and deshaping
config['use_skinning'] = False # and deshaping
config['remove_global_orientation'] = True
config['use_curvature'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
@@ -19,8 +19,8 @@ config['only_curvatures'] = False
ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
assert ablation_count == 1 or ablation_count == 0
config['gpu'] = '1'
config['name'] = 'al_all_post'
config['gpu'] = '5'
config['name'] = 'iccv_weights_test_3'
config['notes'] = ''
config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv
@@ -53,7 +53,7 @@ if config['is_cape']:
config['smpl_res'] = 3638
config['pretrained_model'] = ''
config['pretrained_model'] = ''#'checkpoints/post/iccv_weights_test_2'
###############
config['lr'] = 0.0001# 0.0001
@@ -86,7 +86,7 @@ config['use_one_hot'] = False
config['cvae'] = dict()
config['cvae']['is_label'] = False
config['cvae']['add_noise'] = True
config['cvae']['noise_amount'] = 0.001
config['cvae']['noise_amount'] = 0.00001
config['cvae']['kernel_size'] = 7
config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3]
convnext/models/convnext_tf.py
+15 -2
View File
@@ -2,7 +2,7 @@ from unittest import skip
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
GlobalAveragePooling2D, Layer,
GlobalAveragePooling2D, Layer, Activation, Lambda,
LayerNormalization, Flatten, Reshape, UpSampling2D)
def drop_path(inputs, drop_prob, is_training):
@@ -257,12 +257,25 @@ class ConvNeXtReverse(tf.keras.Model):
self.last = Dense(3, activation='tanh')
self.weights_model = tf.keras.Sequential([Block(dim=dims[-1]*2, kernel_size=kernel_size),
Block(dim=dims[-1]*2, kernel_size=kernel_size),
Block(dim=dims[-1]*2, kernel_size=kernel_size),
Flatten(),
Dense(128 * 24),
Reshape((24, 128)),
Dense(14475),
Lambda(lambda x: tf.transpose(x, (0, 2, 1))),
Lambda(lambda x: tf.nn.softmax(x, axis=-1))])
def call(self, x, training):
weights = self.weights_model(x, training=training)
for i in reversed(range(len(self.upsample_layers))):
x = self.upsample_layers[i](x, training=training)
x = self.stages[i](x, training=training)
x = self.last(x, training=training)
return x
return x, weights
class ConvNeXtReverseWithSkip(tf.keras.Model):
custom_utils/custom_utils.py
+30
View File
@@ -2,6 +2,7 @@ import json
import numpy as np
import tensorflow as tf
import scipy as sp
import cmapy as cm
import plotly.graph_objects as go
from plotly.subplots import make_subplots
@@ -216,6 +217,35 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
return fig
def plot_weights(weights, body, faces):
fig = go.Figure()
colors = np.array([cm.color('jet', 1 - w) for w in weights.numpy()])
# if flip_xy:
# mesh_gt = flip_axis(mesh_gt)
# mesh_pred_unposed = flip_axis(mesh_pred_unposed)
# collision = flip_axis(collision, ind=14475)
fig.add_trace(go.Mesh3d(
x=body[:,0],
y=body[:,2],
z=body[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
vertexcolor=colors
))
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),)
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
fig.update_layout(scene = helper,)
return fig
def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False):
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
custom_utils/smpl.py
+3 -3
View File
@@ -406,7 +406,8 @@ class SMPL:
return T_out, J_out
def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
def run_garment_with_weights(self, pose, betas, v_cloth, weights, J=None):
R = self.run_rodrigues(pose)
if J is None:
J = self.get_joints_template(betas)
@@ -430,13 +431,12 @@ class SMPL:
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
T = tf.tensordot(weights, results, axes=((1), (0)))
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
v = tf.reshape(v, (-1, 4))[:, :3]
return v, tf.squeeze(J)
# gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
cvae.py
+33 -22
View File
@@ -38,7 +38,7 @@ else:
from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, plot_weights
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAENew as CVAE
@@ -71,9 +71,15 @@ def kl_loss(mean, logvar):
return kl_loss
def blend_weights_reg(pred_w):
orig = smpl.weights
l1 = tf.reduce_mean(tf.abs(pred_w - orig[None, 0, :, :]))
return l1 * 1e-3
# 3d loss on posed !!!!!!!!!!!!
# @tf.function
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
@@ -82,7 +88,7 @@ def train_step(batch):
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
mesh_pred = uv2mesh(pred)
@@ -94,25 +100,26 @@ def train_step(batch):
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred, weights)
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
weights_l = blend_weights_reg(weights)
reconstr_loss = l2_loss
reconstr_loss = l2_mesh_loss + weights_l
if config['use_extra_losses']:
if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
reconstr_loss = reconstr_loss + (e_loss) * config['reg_loss_coeff']
else:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss'] and step > 2000:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
@@ -128,8 +135,8 @@ def train_step(batch):
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
# 'train/ms_ssim': ms_ssim,
'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss, 'train/weights_loss': weights_l,
# 'train/ms_ssim': ms_ssim, 'train/l2_loss': l2_loss,
}
@@ -140,7 +147,7 @@ def valid_step(batch):
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
@@ -152,22 +159,23 @@ def valid_step(batch):
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
else:
mesh_pred_posed = mesh_pred
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred, weights)
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
# conc_loss = concept_loss(pred, batch['uv_unposed'])
kl = kl_loss(mean, logvar)
weights_l = blend_weights_reg(weights)
reconstr_loss = l2_loss
reconstr_loss = l2_mesh_loss + weights_l
if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
if config['use_norm_loss']:
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
@@ -182,7 +190,7 @@ def valid_step(batch):
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric,
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar, 'valid/weights_loss': weights_l,
'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim,
}
@@ -194,7 +202,7 @@ def test_step(batch):
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
@@ -208,7 +216,7 @@ def test_step(batch):
mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
else:
mesh_pred_posed = mesh_pred
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred, weights)
mesh_orig_posed = mesh_orig
mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
@@ -250,7 +258,7 @@ def plot_everything(batch: dict, type: str):
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred).numpy()
@@ -277,7 +285,7 @@ def plot_everything(batch: dict, type: str):
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
else:
mesh_posed = mesh_pred
mesh_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred, weights).numpy()
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
@@ -296,13 +304,16 @@ def plot_everything(batch: dict, type: str):
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
# figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], garment_posed[i], faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
figs_weights = { type + '/weights_' + str(i): plot_weights(weights[0,:,i], body[0], F) for i in range(24) }
# figs_weights = { type + '/weights_' + str(i): plot_weights(smpl.weights[0,:,i], body[0], F) for i in range(24) }
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
@@ -476,7 +487,7 @@ def test():
if __name__ == '__main__':
# init_wandb()
init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
models/cvae_model.py
+7 -7
View File
@@ -193,12 +193,12 @@ class CVAENew(tf.keras.Model):
self.noise_amount = config['cvae']['noise_amount']
self.use_curvature = config['use_curvature']
self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
self.use_offsets_over_body = not config['use_skinning']
self.use_offsets_over_body = False
if self.use_offsets_over_body:
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
if self.avg_garment is not None:
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
# if self.avg_garment is not None:
# assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
@@ -213,7 +213,7 @@ class CVAENew(tf.keras.Model):
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z, training=False)
x, weights = self.decoder(z, training=False)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
@@ -223,7 +223,7 @@ class CVAENew(tf.keras.Model):
x = x * uv_mask[:,:,:,tf.newaxis]
return x
return x, weights
def reparameterize(self, mean, logvar, training, fix_noise=False):
@@ -271,7 +271,7 @@ class CVAENew(tf.keras.Model):
z = tf.concat([z, z_cond], axis=-1)
x = self.decoder(z, training=training)
x, weights = self.decoder(z, training=training)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
@@ -281,5 +281,5 @@ class CVAENew(tf.keras.Model):
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar
return x, mean, logvar, weights