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_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
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'])