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

ae updates

Browse Source
This commit is contained in:
luni
2022-10-19 18:42:59 +00:00
parent 8eaf090e4c
commit 1f85598db5
3 changed files with 27 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+14 -14
View File
@@ -33,7 +33,7 @@ from models.cond_encoder import AutoEncoder
def init_wandb(): def init_wandb():
run = wandb.init(project='ae', name=config['name'], notes=config['notes']) run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py') artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
@@ -45,7 +45,7 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py') artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact) run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256) uv2mesh = UV2Mesh(res=256, is_smpl_orig=config['is_cape'])
@tf.function @tf.function
@@ -56,9 +56,9 @@ def train_step(batch):
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed'] mesh_orig = batch['mesh_body_posed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 14475)))) e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 3638 if config['is_cape'] else 14475))))
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 2, batch['uv_body_posed' + resolution_modifier] + 2, max_val=4, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
loss = 1e3 * (1 - ms_ssim) loss = 1e3 * (1 - ms_ssim)
@@ -76,9 +76,9 @@ def valid_step(batch):
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed'] mesh_orig = batch['mesh_body_posed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 14475)))) e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 3638 if config['is_cape'] else 14475))))
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 2, batch['uv_body_posed' + resolution_modifier] + 2, max_val=4, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
loss = 1e3 * (1 - ms_ssim) loss = 1e3 * (1 - ms_ssim)
@@ -99,14 +99,14 @@ def plot_everything(batch: dict, type: str):
img = tf.concat((orig_img, pred_img), axis=2) img = tf.concat((orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
F = np.load("custom_utils/smpl_faces.npy") F = np.load("custom_utils/smpl_faces" + ("_orig" if config['is_cape'] else "") + ".npy")
garment_unposed, faces = mesh_pred.numpy(), F garment_unposed, faces = mesh_pred.numpy(), F
gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])] gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])]
body = uv2mesh(batch['uv_body_posed' + resolution_modifier]).numpy() body = uv2mesh(batch['uv_body_posed' + resolution_modifier]).numpy()
garment_posed = garment_unposed garment_posed = garment_unposed
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces + 14475), axis=0) for i in range(len(garment_posed))] collision_faces = [np.concatenate((F, faces + (3638 if config['is_cape'] else 14475)), axis=0) for i in range(len(garment_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) } figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
@@ -127,7 +127,7 @@ def train():
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = valid_step(data_valid) valid_loss = valid_step(data_valid)
to_log = { **train_loss, **valid_loss } to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 500 == 0): if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
train_figs = plot_everything(batch, 'train') train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid') valid_figs = plot_everything(data_valid, 'valid')
@@ -159,15 +159,15 @@ if __name__ == '__main__':
if config['resolution'] != 512: if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution']) resolution_modifier = '_' + str(config['resolution'])
uv_mask_body = np.load('custom_utils/mask_256x512.npy') uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
factory = DatasetFactory() factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
model = AutoEncoder(config['cond_ae']) model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr']) optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
auto_encoder_curv.py
+6 -6
View File
@@ -35,7 +35,7 @@ from models.cond_encoder import AutoEncoder
testing = False testing = False
def init_wandb(): def init_wandb():
if not testing: if not testing:
run = wandb.init(project='ae', name=config['name'], notes=config['notes']) run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/')
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py') artifact.add_file('auto_encoder_curv.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
@@ -47,7 +47,7 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py') artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact) run.log_artifact(artifact)
uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt']) uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7) g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
@@ -202,10 +202,10 @@ if __name__ == '__main__':
factory = DatasetFactory() factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
model = AutoEncoder(config['cond_ae']) model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr']) optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
auto_encoder_static.py
+7 -9
View File
@@ -35,10 +35,8 @@ from custom_utils.custom_utils import batch_to_dict, plot_3d
from models.cond_encoder import AutoEncoder from models.cond_encoder import AutoEncoder
def init_wandb(): def init_wandb():
run = wandb.init(project='ae', name=config['name'], run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
notes=config['notes'])
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_static.py') artifact.add_file('auto_encoder_static.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
@@ -50,7 +48,7 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py') artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact) run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt']) uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -102,7 +100,7 @@ def plot_everything(batch: dict, type: str):
img = tf.concat((orig_img, pred_img), axis=2) img = tf.concat((orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy") F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool) mask = batch['mesh_mask'].numpy().astype(bool)
garment_static, faces = maskFacesBatch(mesh_pred.numpy(), F, mask) garment_static, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
@@ -200,10 +198,10 @@ if __name__ == '__main__':
factory = DatasetFactory() factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
model : AutoEncoder = AutoEncoder(config['cond_ae']) model : AutoEncoder = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr']) optimizer = tf.optimizers.Adam(learning_rate=config['lr'])