diff --git a/auto_encoder_body.py b/auto_encoder_body.py index fbd2a13..5473289 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -33,7 +33,7 @@ from models.cond_encoder import AutoEncoder 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.add_file('auto_encoder_body.py') artifact.add_file('losses.py') @@ -45,7 +45,7 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) -uv2mesh = UV2Mesh(res=256) +uv2mesh = UV2Mesh(res=256, is_smpl_orig=config['is_cape']) @tf.function @@ -56,9 +56,9 @@ def train_step(batch): mesh_pred = uv2mesh(pred) 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)) - 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) @@ -76,9 +76,9 @@ def valid_step(batch): mesh_pred = uv2mesh(pred) 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)) - 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) @@ -99,14 +99,14 @@ def plot_everything(batch: dict, type: str): img = tf.concat((orig_img, pred_img), axis=2) 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 gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])] body = uv2mesh(batch['uv_body_posed' + resolution_modifier]).numpy() garment_posed = garment_unposed 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))] 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'])) 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): train_figs = plot_everything(batch, 'train') valid_figs = plot_everything(data_valid, 'valid') @@ -159,15 +159,15 @@ if __name__ == '__main__': if config['resolution'] != 512: 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) 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']) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) - 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']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) + 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, 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'], 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, is_cape=config['is_cape']) model = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index 15b629e..33ff70b 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -35,7 +35,7 @@ from models.cond_encoder import AutoEncoder testing = False def init_wandb(): 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.add_file('auto_encoder_curv.py') artifact.add_file('losses.py') @@ -47,7 +47,7 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') 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) @@ -202,10 +202,10 @@ if __name__ == '__main__': 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']) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) - 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']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) + 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, 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'], 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, is_cape=config['is_cape']) model = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) diff --git a/auto_encoder_static.py b/auto_encoder_static.py index a6044a7..d1244a2 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -35,10 +35,8 @@ from custom_utils.custom_utils import batch_to_dict, plot_3d from models.cond_encoder import AutoEncoder - 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.add_file('auto_encoder_static.py') artifact.add_file('losses.py') @@ -50,7 +48,7 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') 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) 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) garment_static, faces = maskFacesBatch(mesh_pred.numpy(), F, mask) gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] @@ -200,10 +198,10 @@ if __name__ == '__main__': 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']) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) - 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']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) + 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, 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'], 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, is_cape=config['is_cape']) model : AutoEncoder = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr'])