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

working auto encoders

Browse Source
This commit is contained in:
luni
2022-09-25 22:12:36 +00:00
parent 26916c00e7
commit 040b29e898
9 changed files with 41 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+10 -7
View File
@@ -31,7 +31,7 @@ def init_wandb():
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config.py')
artifact.add_file('configs/config_base.py')
artifact.add_file('configs/config_ae_body.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
@@ -45,7 +45,7 @@ uv2mesh = UV2Mesh(res=256)
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=True)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
@@ -87,7 +87,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=False)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
@@ -133,7 +133,8 @@ def train():
wandb.log({ **train_loss, **valid_loss }, step=step)
if step % 10000 == 0:
manager.save(checkpoint_number=step)
manager_enc.save(checkpoint_number=step)
manager_dec.save(checkpoint_number=step)
print(step)
@@ -141,7 +142,7 @@ def train():
if __name__ == '__main__':
# init_wandb()
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -165,7 +166,9 @@ if __name__ == '__main__':
model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3)
ckpt_dec = tf.train.Checkpoint(model=model.decoder)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
train()
auto_encoder_curv.py
+15 -52
View File
@@ -1,11 +1,12 @@
import os
from configs.config_ae_curv import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import time
import argparse
import wandb
import random
import numpy as np
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
@@ -16,7 +17,6 @@ from tf_utils import check_tf_version, allow_memory_growth
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from models.cond_encoder import AutoEncoder
from configs.config_ae_curv import config
testing = False
@@ -26,7 +26,7 @@ def init_wandb():
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config.py')
artifact.add_file('configs/config_base.py')
artifact.add_file('configs/config_ae_curv.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
@@ -37,45 +37,6 @@ def init_wandb():
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred, faces):
mesh_pred[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred[:,0],
y=mesh_pred[:,1],
z=mesh_pred[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
def get_mean_iou(y_true, y_pred):
g_mean_iou.reset_state()
@@ -112,13 +73,13 @@ 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:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred = model([uv_curvature, batch['uv_mask']])
pred = model([uv_curvature, batch['uv_mask']], training=True)
ce_uv_loss = scce(batch['uv_curvature'], pred)
@@ -166,20 +127,21 @@ def train():
to_log = { **train_loss, **valid_loss}
if step % 2500 == 0:
uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7)
pred = model([uv_curvature, data_valid['uv_mask']])
pred = model([uv_curvature, data_valid['uv_mask']], training=False)
orig_img = data_valid['uv_curvature']
pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
img = tf.concat((orig_img, pred_img), axis=2).numpy()
# imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] }
to_log = { **to_log, **imgs}
wandb.log(to_log, step=step)
print(step)
if step % 2500 == 0:
manager.save(checkpoint_number=step)
if step % 10000 == 0:
manager_enc.save(checkpoint_number=step)
manager_dec.save(checkpoint_number=step)
step += 1
@@ -201,7 +163,7 @@ def test():
if __name__ == '__main__':
# init_wandb()
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -223,9 +185,10 @@ if __name__ == '__main__':
model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3)
ckpt_dec = tf.train.Checkpoint(model=model.decoder)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
if testing:
test()
auto_encoder_static.py
+4 -4
View File
@@ -83,7 +83,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_static'], batch['uv_mask']])
pred = model([batch['uv_static'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
@@ -151,7 +151,7 @@ def test():
if __name__ == '__main__':
# init_wandb()
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -178,5 +178,5 @@ if __name__ == '__main__':
ckpt_dec = tf.train.Checkpoint(model=model.decoder)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
test()
# train()
# test()
train()
configs/config_ae_body.py
+1 -1
View File
@@ -1,6 +1,6 @@
from configs.config_base import config
config['gpu'] = '1'
config['gpu'] = '0'
config['lr'] = 0.0001
config['name'] = 'ae_body_9'
configs/config_ae_curv.py
+2 -1
View File
@@ -1,7 +1,8 @@
from configs.config_base import config
config['gpu'] = '2'
config['lr'] = 0.001
config['name'] = 'ae_body_9'
config['name'] = 'ae_curv_9'
config['cond_ae']['is_label'] = True
configs/config_ae_static.py
+2 -2
View File
@@ -1,8 +1,8 @@
from configs.config_base import config
config['gpu'] = '0'
config['gpu'] = '1'
config['lr'] = 0.0005
config['name'] = 'ae_static_7_x'
config['name'] = 'ae_static_7_2'
config['signature_dynamic'] = { 'id': (), }
config['signature_static'] = { 'uv_static': (512, 256, 3),
dataset.py
+2 -2
View File
@@ -54,7 +54,7 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
class StaticDataset():
def __init__(self, folder, signature, use_filled=True, resolutions=[512]) -> None:
def __init__(self, folder, signature, use_filled=False, resolutions=[512]) -> None:
self.store = dict()
self.signature = signature
self.folder = folder
@@ -139,7 +139,7 @@ class DatasetFactory():
return tf_signature
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True, resolutions=[512]):
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512]):
# assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 128, 512]))
models/base.py
+4 -1
View File
@@ -64,7 +64,10 @@ class AEEncoder(tf.keras.Model):
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
return x
if len(inputs) == 2:
return x, inputs[1]
else:
return x
class Decoder(tf.keras.Model):
models/cond_encoder.py
+1 -1
View File
@@ -8,7 +8,7 @@ class AutoEncoder(tf.keras.Model):
super(AutoEncoder, self).__init__()
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
self.decoder : Decoder= Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
def call(self, inputs, training=None):