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

making ready for skirts

Browse Source
This commit is contained in:
luni
2022-09-30 21:53:17 +00:00
parent 6a2f9ff4ff
commit dca5368704
22 changed files with 700 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+13 -11
View File
@@ -26,10 +26,15 @@ import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
if config['is_skirt']:
from losses import edge_loss_skirt as edge_loss
else:
from losses import edge_loss
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from losses import edge_loss, uv_mask_loss, bend_loss, kl_loss
from losses import uv_mask_loss, bend_loss, kl_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
@@ -52,7 +57,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_skirt=config['is_skirt'])
def log_normal_pdf(sample, mean, logvar, raxis=1):
@@ -72,10 +77,7 @@ def kl_loss(mean, logvar):
def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']],
training=True, use_mask=True)
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
@@ -94,7 +96,7 @@ def train_step(batch):
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
reconstr_loss = l2_loss + e_loss + c_loss
reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
reconstr_loss = tf.reduce_mean(reconstr_loss)
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
@@ -152,8 +154,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=False, use_mask=True)
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
mesh_pred = uv2mesh(pred).numpy()
mesh_orig = batch['mesh_unposed']
@@ -169,7 +170,7 @@ def plot_everything(batch: dict, type: str):
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
F = np.load("custom_utils/smpl_faces.npy")
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
@@ -198,7 +199,7 @@ def sample(folder):
load_model(model,folder)
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2
@@ -322,6 +323,7 @@ if __name__ == '__main__':
args = vars(parser.parse_args())
smpl = SMPL()
skirt_modifier = '_skirt' if config['is_skirt'] else ''
factory = DatasetFactory()