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

train encoder on raw curvature values

Browse Source
This commit is contained in:
luni
2023-01-27 11:13:10 +00:00
parent a3c5b526b8
commit 6e6b0a6604
8 changed files with 60 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_curv.py
+22 -29
View File
@@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
testing = False
def init_wandb():
if not testing:
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/')
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py')
artifact.add_file('losses.py')
@@ -47,7 +47,9 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -100,22 +102,18 @@ def train_step(batch):
with tf.GradientTape() as tape:
pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
ce_uv_loss = tf.reduce_mean(tf.abs(batch['uv_curvature'] - pred))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred, batch['uv_curvature'], max_val=1))
mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
ce_mesh_loss = scce(mesh_orig, mesh_pred)
ce_mesh_loss = tf.reduce_mean(tf.abs(batch['mesh_curvature'] - mesh_pred[:,:,0]))
loss = ce_uv_loss + ce_mesh_loss
mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, batch['uv_curvature'])
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou,
'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss }
return { 'train/loss': loss, 'train/uv_l1': ce_uv_loss, 'train/ssim': ms_ssim,
'train/mesh_l1': ce_mesh_loss }
@@ -123,35 +121,30 @@ def train_step(batch):
def valid_step(batch):
pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
ce_uv_loss = tf.reduce_mean(tf.abs(batch['uv_curvature'] - pred))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred, batch['uv_curvature'], max_val=1))
mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
ce_mesh_loss = scce(mesh_orig, mesh_pred)
ce_mesh_loss = tf.reduce_mean(tf.abs(batch['mesh_curvature'] - mesh_pred[:,:,0]))
loss = ce_uv_loss + ce_mesh_loss
mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, batch['uv_curvature'])
return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
return { 'valid/loss': loss, 'valid/uv_l1': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mesh_l1': ce_mesh_loss }
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
orig_img = tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.dtypes.int32)
pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
orig_img = batch['uv_curvature']
pred_img = pred
img = tf.concat((orig_img, pred_img), axis=2).numpy()
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], img.shape[2], 3), np.uint8)
for i in range(7):
color_img[img == i, :] = colors[i]
color_img = [wandb.Image(color_img[i]) for i in range(4)]
# 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], img.shape[2], 3), np.uint8)
# for i in range(7):
# color_img[img == i, :] = colors[i]
# color_img = [wandb.Image(color_img[i]) for i in range(4)]
imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])], type + '/color_images': color_img }
imgs = { type + '/images': [wandb.Image((img[i,:,:,np.newaxis]-img[i,:,:,np.newaxis].min()) / (img[i,:,:,np.newaxis].max()-img[i,:,:,np.newaxis].min())) for i in range(img.shape[0])], } # type + '/color_images': color_img
return imgs
@@ -206,7 +199,7 @@ def test():
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
# allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())