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

backup

Browse Source
This commit is contained in:
luni
2022-09-02 14:15:54 +00:00
parent ee8f1cb68c
commit 977e49900d
16 changed files with 2241 additions and 272 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_curv.py
+16 -10
View File
@@ -5,6 +5,9 @@ import wandb
import random
import numpy as np
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
@@ -24,14 +27,15 @@ from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='ae_curvature_6'
name='ae_curvature_7'
testing = False
def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
if not testing:
run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128')
run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('stylegan2_ae/layers/mapping_block.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
@@ -282,19 +286,21 @@ def train():
train_loss = train_step(batch)
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = valid_step(data_valid)
# data_valid = valid_iterator.get_next()
# data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = {}#valid_step(data_valid)
data_valid = batch
to_log = { **train_loss, **valid_loss}
if step % 1000 == 0:
if step % 2500 == 0:
uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7)
pred = model([uv_curvature, data_valid['uv_mask']])
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 = { '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)] }
to_log = { **to_log, **imgs}
wandb.log(to_log, step=step)
@@ -351,10 +357,10 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=8,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
batch_size=8,