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

lot of things

Browse Source
This commit is contained in:
luni
2022-08-06 12:27:24 +00:00
parent cd498e79f3
commit ee8f1cb68c
28 changed files with 4595 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -142,4 +142,5 @@ models/
models_archive/
temp.jpg
corrupt/
wandb copy/
wandb copy/
checkpoints/
auto_encoder.py
+280 -86
View File
@@ -1,7 +1,6 @@
import os
import time
import argparse
from cv2 import polarToCart
import wandb
import random
import numpy as np
@@ -9,6 +8,7 @@ import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
@@ -16,16 +16,25 @@ from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_unposed_1', notes='convnext + stylegan generator')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder.py')
artifact.add_file('stylegan2_ae/layers/mapping_block.py')
run.log_artifact(artifact)
name='ae_body_1'
# os.environ["CUDA_VISIBLE_DEVICES"] = "6"
# run = wandb.init(project='stylegan', name=name,
# notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
# fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
# . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.')
# artifact = wandb.Artifact('code', type='code')
# artifact.add_file('auto_encoder.py')
# artifact.add_file('stylegan2_ae/layers/mapping_block.py')
# artifact.add_file('convnext/models/convnext_tf.py')
# run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
@@ -67,8 +76,46 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
return filtered_resolutions, filtered_featuremaps
def plot_3d(mesh_gt, mesh_pred, faces):
mesh_pred[:,1] += 2
# def plot_3d(mesh_gt, mesh_pred, faces):
# mesh_pred[:,1] += 2
# fig = go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_gt[:,0],
# y=mesh_gt[:,1],
# z=mesh_gt[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_pred[:,0],
# y=mesh_pred[:,1],
# z=mesh_pred[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2],
# )
# ])
# fig.update_layout(scene_aspectmode='manual',
# scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(
# scene = dict(
# xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,3],),
# zaxis = dict(nticks=4, range=[-1,1],),)
# )
# return fig
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
@@ -82,13 +129,24 @@ def plot_3d(mesh_gt, mesh_pred, faces):
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred[:,0],
y=mesh_pred[:,1],
z=mesh_pred[:,2],
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
@@ -103,31 +161,35 @@ def plot_3d(mesh_gt, mesh_pred, faces):
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x, training=training)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
g_featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, 512)
g_params = {
'z_dim': 256,
'w_dim': 256,
'labels_dim': -1,
'n_mapping': 8,
'resolutions': train_resolutions,
'featuremaps': train_g_featuremaps,
}
# test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
self.decoder = Generator(g_params)
# _ = generator(test_latent)
bottleneck_size = 32
self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, g_params['z_dim']], include_top=False)
self.pooling = tf.keras.layers.GlobalAveragePooling2D()
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size)
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
def call(self, inputs, training=None):
@@ -135,105 +197,237 @@ class AutoEncoder(tf.keras.Model):
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.pooling(x)
# x = self.pooling(x)
# x = self.reshape(x)
# x = self.flatten(x)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
x = x * uv_mask[:,tf.newaxis,:,:]
return x
def draw_cicle(shape,diamiter):
'''
Input:
shape : tuple (height, width)
diameter : scalar
Output:
np.array of shape that says True within a circle with diamiter = around center
'''
assert len(shape) == 2
TF = np.zeros(shape,dtype=bool)
center = np.array(TF.shape)/2.0
for iy in range(shape[0]):
for ix in range(shape[1]):
TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
return (TF)
TFcircleIN = draw_cicle(shape=(256,256), diamiter=15)
TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
TFcircleOUT = ~TFcircleIN
def fft_loss_per_channel(pred, gt):
pred = tf.complex(pred, tf.zeros_like(pred))
gt = tf.complex(gt, tf.zeros_like(gt))
fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
return total, low_loss, high_loss
# ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
# img = tf.signal.fftshift(fft_gt[0][0])
# ex.imshow(tf.math.abs(tf.math.log(img))).show()
# low = tf.where(TFcircleIN, img, tf.zeros_like(img))
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
# high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
# # ex.imshow(tf.math.abs(tf.math.log(high))).show()
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
print()
@tf.function
def train_step(batch):
uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
with tf.GradientTape() as tape:
pred = model([uv_unposed, uv_mask])
pred = model([batch['uv_unposed'], batch['uv_mask']])
# pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred = uv2mesh(pred)
mesh_orig = mesh_unposed
# mesh_pred = pred
mesh_orig = batch['mesh_unposed']
e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask)
c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask)
loss = l1_loss + e_loss + c_loss
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
# l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2))
rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
# loss = fft_loss_pc # + e_loss + c_loss
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
# mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric }
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
@tf.function
def valid_step(batch):
uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
pred = model([batch['uv_unposed'], batch['uv_mask']])
pred = model([uv_unposed, uv_mask])
mesh_pred = uv2mesh(pred)
mesh_orig = mesh_unposed
mesh_orig = batch['mesh_unposed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2))
fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
}
e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask)
c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask)
loss = l1_loss + e_loss + c_loss
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric }
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_unposed'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
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")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + 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[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **{ type + '/images': imgs } }
if __name__ == '__main__':
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
name = 'ae_unposed_1'
dataset_train = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(8).repeat()
dataset_valid = DatasetSimple('/home/workspace/data/dataset/stylegan_new/valid').cache().batch(8).repeat()
smpl = SMPL()
# R = smpl.run_rodrigues((np.zeros((8,72), dtype=np.float32)))
# zero_body = smpl.run(R, np.zeros((8, 11), dtype=np.float32))[0]
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=8,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
batch_size=8,
signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
signature=signature_static)
model = AutoEncoder()
# optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
optimizer = tf.optimizers.Adam()
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
folder = 'models/ae_logs/'
train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
# train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
# valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
valid_iterator = iter(dataset_valid)
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
train_loss = train_step(batch)
for key in train_loss.keys():
with train_summary_writer.as_default():
tf.summary.scalar(key, train_loss[key], step=step)
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
# data_valid = iter(dataset_valid).get_next()
# valid_loss = valid_step(data_valid)
# for key in valid_loss.keys():
# with valid_summary_writer.as_default():
# tf.summary.scalar(key, valid_loss[key], step=step)
if step % 250 == 0:
pred = model([data_valid[0], data_valid[2]])
if step % 1000 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
mesh_pred = uv2mesh(pred)
mesh_orig = data_valid[1]
orig_img = tf.transpose(data_valid[0], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
img = tf.concat((orig_img, pred_img), axis=2)
with train_summary_writer.as_default():
tf.summary.image('images', img, step=step)
F = np.load("custom_utils/smpl_faces.npy")
garment, faces = maskFacesBatch(mesh_pred.numpy(), F, data_valid[3].numpy().astype(bool))
gt = [mesh_orig[i].numpy()[data_valid[3][i].numpy().astype(bool), :] for i in range(pred.shape[0])]
for i in range(pred.shape[0]):
fig = plot_3d(gt[i], garment[i], faces[i])
wandb.log({"garment_" + str(i): fig})
print(step)
auto_encoder_body.py
+435
View File
@@ -0,0 +1,435 @@
import os
import time
import argparse
import wandb
import random
import cv2
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='ae_body_5'
def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
run = wandb.init(project='vae', name=name,
notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
. lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
two more downscaling than up, so decoder output is resized wto final resolution\
larger batch_size.\
flattened bottleneck with 128 latent_dim\ remove global orientation')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
class DatasetSimple(tf.data.Dataset):
def _generator(folder):
# Opening the file
folder = folder.decode("utf-8")
uv_folder = os.path.join(folder, 'uv_posed')
frames = os.listdir(uv_folder)
random.shuffle(frames)
for frame in frames:
id, frame_nr = frame.split('.')[0].split('_')
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
args=(folder, )
)
def filter_resolutions_featuremaps(resolutions, featuremaps, res):
index = resolutions.index(res)
filtered_resolutions = resolutions[:index + 1]
filtered_featuremaps = featuremaps[:index + 1]
return filtered_resolutions, filtered_featuremaps
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[3, 3, 9, 9, 3, 3], dims=[32, 64, 128, 256, 256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x, _ = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# x = self.pooling(x)
# x = self.reshape(x)
# x = self.flatten(x)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]
return x
def draw_cicle(shape,diamiter):
'''
Input:
shape : tuple (height, width)
diameter : scalar
Output:
np.array of shape that says True within a circle with diamiter = around center
'''
assert len(shape) == 2
TF = np.zeros(shape,dtype=bool)
center = np.array(TF.shape)/2.0
for iy in range(shape[0]):
for ix in range(shape[1]):
TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
return (TF)
TFcircleIN = draw_cicle(shape=(256,256), diamiter=15)
TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
TFcircleOUT = ~TFcircleIN
def fft_loss_per_channel(pred, gt):
pred = tf.complex(pred, tf.zeros_like(pred))
gt = tf.complex(gt, tf.zeros_like(gt))
fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
return total, low_loss, high_loss
# ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
# img = tf.signal.fftshift(fft_gt[0][0])
# ex.imshow(tf.math.abs(tf.math.log(img))).show()
# low = tf.where(TFcircleIN, img, tf.zeros_like(img))
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
# high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
# # ex.imshow(tf.math.abs(tf.math.log(high))).show()
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
print()
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred = model([batch['uv_body_posed'], batch['uv_mask']])
# pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred = uv2mesh(pred)
# mesh_pred = pred
mesh_orig = batch['mesh_body_posed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
# l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_body_posed'] - pred)), axis=1) * batch['uv_mask'])
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 1, max_val=2))
rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
# loss = fft_loss_pc # + e_loss + c_loss
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_body_posed'])
# mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
@tf.function
def valid_step(batch):
pred = model([batch['uv_body_posed'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_body_posed'] - pred)), axis=1) * batch['uv_mask'])
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 1, max_val=2))
fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_body_posed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
}
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_body_posed'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
orig_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
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")
mask = batch['mesh_mask'].numpy().astype(bool)
# garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
garment_unposed, faces = mesh_pred.numpy(), F
# gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
# mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
# garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
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[i] + 14475), 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_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[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **{ type + '/images': imgs } }
def train():
valid_iterator = iter(dataset_valid)
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
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)
if step % 1000 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = {}#plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
if step % 2500 == 0:
manager.save(checkpoint_number=step)
print(step)
step += 1
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
uv_mask_body = np.load('custom_utils/mask_256x512.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()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=32, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model = AutoEncoder()
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=0.0001)
folder = 'models/ae_logs/'
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
train()
auto_encoder_curv.py
+379
View File
@@ -0,0 +1,379 @@
import os
import time
import argparse
import wandb
import random
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='ae_curvature_6'
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')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py')
artifact.add_file('stylegan2_ae/layers/mapping_block.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
class DatasetSimple(tf.data.Dataset):
def _generator(folder):
# Opening the file
folder = folder.decode("utf-8")
uv_folder = os.path.join(folder, 'uv_posed')
frames = os.listdir(uv_folder)
random.shuffle(frames)
for frame in frames:
id, frame_nr = frame.split('.')[0].split('_')
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
args=(folder, )
)
def filter_resolutions_featuremaps(resolutions, featuremaps, res):
index = resolutions.index(res)
filtered_resolutions = resolutions[:index + 1]
filtered_featuremaps = featuremaps[:index + 1]
return filtered_resolutions, filtered_featuremaps
def plot_3d(mesh_gt, mesh_pred, faces):
mesh_pred[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred[:,0],
y=mesh_pred[:,1],
z=mesh_pred[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
latent_dim = 128
self.decoder = Decoder(bottleneck_size, latent_dim)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
def call(self, inputs, training=None):
x, uv_mask = inputs
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
def get_mean_iou(y_true, y_pred):
g_mean_iou.reset_state()
g_mean_iou.update_state(y_true, y_pred)
return g_mean_iou.result()
# https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py
def dice_dense(prediction, ground_truth, weight_map=None):
"""
Computing mean-class Dice similarity.
:param prediction: last dimension should have ``num_classes``
:param ground_truth: segmentation ground truth (encoded as a binary matrix)
last dimension should be ``num_classes``
:param weight_map:
:return: ``1.0 - mean(Dice similarity per class)``
"""
if weight_map is not None:
raise NotImplementedError
prediction = tf.cast(prediction, dtype=tf.float32)
ground_truth = tf.cast(ground_truth, dtype=tf.float32)
ground_truth = tf.reshape(ground_truth, prediction.shape)
# computing Dice over the spatial dimensions
reduce_axes = list(range(len(prediction.shape) - 1))
dice_numerator = 2.0 * tf.reduce_sum(
prediction * ground_truth, axis=reduce_axes)
dice_denominator = \
tf.reduce_sum(tf.square(prediction), axis=reduce_axes) + \
tf.reduce_sum(tf.square(ground_truth), axis=reduce_axes)
epsilon = 0.00001
dice_score = (dice_numerator + epsilon) / (dice_denominator + epsilon)
return tf.reduce_mean(dice_score)
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred = model([uv_curvature, batch['uv_mask']])
ce_uv_loss = scce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], batch['uv_curvature'][:,:,:,tf.newaxis], max_val=7*2))
# loss = 1e3 * (1 - ms_ssim)
loss = ce_uv_loss
# mesh_pred = uv2mesh(tf.tile(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], [1,1,1,3]))[:,:,0]
# mesh_orig = tf.one_hot(batch['mesh_curvature'] + tf.cast(batch['mesh_mask'], tf.int32), 7)
# ce_mesh_loss = scce(mesh_orig, mesh_pred)
mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, uv_curvature)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice #'ce_mesh_loss': ce_mesh_loss
}
@tf.function
def valid_step(batch):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred = model([uv_curvature, batch['uv_mask']])
ce_uv_loss = scce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
loss = ce_uv_loss
mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, uv_curvature)
return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/dice': dice #'ce_mesh_loss': ce_mesh_loss
}
def train():
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
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)
to_log = { **train_loss, **valid_loss}
if step % 1000 == 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)] }
to_log = { **to_log, **imgs}
wandb.log(to_log, step=step)
print(step)
if step % 2500 == 0:
manager.save(checkpoint_number=step)
step += 1
def test():
ckpt.restore(manager.latest_checkpoint)
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred = model([uv_curvature, batch['uv_mask']])
orig_img = batch['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()
ex.imshow(img[0] / 6.0).show()
print()
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
scce = tf.keras.losses.SparseCategoricalCrossentropy()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
'uv_curvature': (512, 256),
'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=8,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
batch_size=8,
signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
signature=signature_static)
valid_iterator = iter(dataset_valid)
model = AutoEncoder()
# optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.001, decay_steps=10000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
folder = 'models/ae_logs/'
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
if testing:
test()
else:
train()
auto_encoder_static.py
+450
View File
@@ -0,0 +1,450 @@
import os
import time
import argparse
import wandb
import random
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='ae_static_3'
def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
run = wandb.init(project='vae', name=name,
notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
. lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
convnext first stage down/upscale by factor fo 2 instead of 4, and 5 stages instead of 4 so bottlneck same size\
flattened latent dim to 128')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_static.py')
artifact.add_file('losses.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
class DatasetSimple(tf.data.Dataset):
def _generator(folder):
# Opening the file
folder = folder.decode("utf-8")
uv_folder = os.path.join(folder, 'uv_posed')
frames = os.listdir(uv_folder)
random.shuffle(frames)
for frame in frames:
id, frame_nr = frame.split('.')[0].split('_')
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
args=(folder, )
)
def filter_resolutions_featuremaps(resolutions, featuremaps, res):
index = resolutions.index(res)
filtered_resolutions = resolutions[:index + 1]
filtered_featuremaps = featuremaps[:index + 1]
return filtered_resolutions, filtered_featuremaps
# def plot_3d(mesh_gt, mesh_pred, faces):
# mesh_pred[:,1] += 2
# fig = go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_gt[:,0],
# y=mesh_gt[:,1],
# z=mesh_gt[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_pred[:,0],
# y=mesh_pred[:,1],
# z=mesh_pred[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2],
# )
# ])
# fig.update_layout(scene_aspectmode='manual',
# scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(
# scene = dict(
# xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,3],),
# zaxis = dict(nticks=4, range=[-1,1],),)
# )
# return fig
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
self.flatten = tf.keras.layers.Flatten()
self.dense = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(8 * 4 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((8, 4, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3, 3])), dims=list(reversed([bottleneck_size, 128, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
latent_dim = 128
self.encoder = ConvNeXt(depths=[3, 3, 9, 3, 3], dims=[32, 64, 128, 128, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x, uv_mask = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# x = self.pooling(x)
# x = self.reshape(x)
# x = self.flatten(x)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
x = x * uv_mask[:,tf.newaxis,:,:]
return x
def draw_cicle(shape,diamiter):
'''
Input:
shape : tuple (height, width)
diameter : scalar
Output:
np.array of shape that says True within a circle with diamiter = around center
'''
assert len(shape) == 2
TF = np.zeros(shape,dtype=bool)
center = np.array(TF.shape)/2.0
for iy in range(shape[0]):
for ix in range(shape[1]):
TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
return (TF)
TFcircleIN = draw_cicle(shape=(256,256), diamiter=15)
TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
TFcircleOUT = ~TFcircleIN
def fft_loss_per_channel(pred, gt):
pred = tf.complex(pred, tf.zeros_like(pred))
gt = tf.complex(gt, tf.zeros_like(gt))
fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
return total, low_loss, high_loss
# ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
# img = tf.signal.fftshift(fft_gt[0][0])
# ex.imshow(tf.math.abs(tf.math.log(img))).show()
# low = tf.where(TFcircleIN, img, tf.zeros_like(img))
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
# high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
# # ex.imshow(tf.math.abs(tf.math.log(high))).show()
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
print()
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred = model([batch['uv_static'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
# l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * batch['uv_mask'])
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_static'])
# mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
@tf.function
def valid_step(batch):
pred = model([batch['uv_static'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * batch['uv_mask'])
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_static'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
}
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_static'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static']
orig_img = tf.transpose(batch['uv_static'], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
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")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = gt
garment_posed = mesh_pred
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), 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[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **{ type + '/images': imgs } }
def train():
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
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)
if step % 1000 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
print(step)
step += 1
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { 'uv_static': (3, 512, 256,),
'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=8, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model = AutoEncoder()
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
folder = 'models/ae_logs/'
valid_iterator = iter(dataset_valid)
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
train()
convnext/models/convnext_tf.py
+305 -19
View File
@@ -1,9 +1,9 @@
from unittest import skip
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import (Conv2D, Dense, DepthwiseConv2D,
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
GlobalAveragePooling2D, Layer,
LayerNormalization)
LayerNormalization, Flatten, Reshape, UpSampling2D)
def drop_path(inputs, drop_prob, is_training):
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
@@ -109,6 +109,299 @@ class ConvNeXt(tf.keras.Model):
LayerNormalization(epsilon=1e-6)]
)
self.downsample_layers.append(stem)
for i in range(len(depths) - 1):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
)
self.downsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(len(depths)):
stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
if self.include_top:
self.avg = GlobalAveragePooling2D()
self.norm = LayerNormalization(epsilon=1e-6) # final norm layer
self.head = Dense(num_classes)
else:
self.avg = None
self.norm = None
self.head = None
def forward_features(self, x):
for i in range(len(self.downsample_layers)):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
return x
def call(self, x):
x = self.forward_features(x)
if self.include_top:
x = self.avg(x)
x = self.norm(x)
x = self.head(x)
return x
class ConvNeXtWithSkip(tf.keras.Model):
r""" ConvNeXt
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
num_classes (int): Number of classes for classification head. Default: 1000
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
include_top (bool): whether to add head or just use it as feature extractor. Default: True
drop_path_rate (float): Stochastic depth rate. Default: 0.
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
"""
def __init__(self,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
drop_path_rate=0., layer_scale_init_value=1e-6,
):
super().__init__()
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
LayerNormalization(epsilon=1e-6)]
)
self.downsample_layers.append(stem)
for i in range(len(depths) - 1):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
)
self.downsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(len(depths)):
stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
def call(self, x):
skip_features = []
for i in range(len(self.downsample_layers)):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
skip_features.append(x)
return x, skip_features
class ConvNeXtReverse(tf.keras.Model):
r""" ConvNeXt
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
num_classes (int): Number of classes for classification head. Default: 1000
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
include_top (bool): whether to add head or just use it as feature extractor. Default: True
drop_path_rate (float): Stochastic depth rate. Default: 0.
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
"""
def __init__(self,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
):
super().__init__()
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
# UpSampling2D((4, 4), interpolation='nearest'),
# Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
# Dense(dims[0]),
LayerNormalization(epsilon=1e-6)]
)
self.upsample_layers.append(stem)
for i in range(len(depths) - 1):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
# UpSampling2D((2,2), interpolation='nearest'),
# Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
# Dense(dims[i+1]),
])
self.upsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(len(depths)):
stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
self.last = Dense(3, activation='tanh')
def call(self, x):
for i in reversed(range(len(self.upsample_layers))):
x = self.upsample_layers[i](x)
x = self.stages[i](x)
x = self.last(x)
return x
class ConvNeXtReverseWithSkip(tf.keras.Model):
r""" ConvNeXt
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
num_classes (int): Number of classes for classification head. Default: 1000
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
include_top (bool): whether to add head or just use it as feature extractor. Default: True
drop_path_rate (float): Stochastic depth rate. Default: 0.
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
"""
def __init__(self,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
):
super().__init__()
assert len(depths) == len(dims)
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
LayerNormalization(epsilon=1e-6)]
)
self.upsample_layers.append(stem)
for i in range(len(depths) - 1):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
)
self.upsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(len(depths)):
stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
self.last = Dense(3, activation='tanh')
def call(self, x):
x, skip_features = x
for i in reversed(range(len(self.upsample_layers))):
x = tf.concat((x, skip_features[i]), axis=-1)
x = self.upsample_layers[i](x)
x = self.stages[i](x)
x = self.last(x)
return x
class ConvNeXtReverseLabels(tf.keras.Model):
r""" ConvNeXt
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
num_classes (int): Number of classes for classification head. Default: 1000
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
include_top (bool): whether to add head or just use it as feature extractor. Default: True
drop_path_rate (float): Stochastic depth rate. Default: 0.
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
"""
def __init__(self,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
nr_of_classes=7):
super().__init__()
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
LayerNormalization(epsilon=1e-6)]
)
self.downsample_layers.append(stem)
for i in range(3):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
)
self.downsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(4):
stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
self.last = Dense(nr_of_classes, activation='softmax')
def call(self, x):
for i in reversed(range(4)):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
x = self.last(x)
return x
class ConvNeXtDiscriminator(tf.keras.Model):
r""" ConvNeXt
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
num_classes (int): Number of classes for classification head. Default: 1000
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
include_top (bool): whether to add head or just use it as feature extractor. Default: True
drop_path_rate (float): Stochastic depth rate. Default: 0.
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
"""
def __init__(self, num_classes=1000,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
):
super().__init__()
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
LayerNormalization(epsilon=1e-6)]
)
self.downsample_layers.append(stem)
for i in range(3):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
@@ -127,28 +420,21 @@ class ConvNeXt(tf.keras.Model):
self.stages.append(stage)
cur += depths[i]
if self.include_top:
self.avg = GlobalAveragePooling2D()
self.norm = LayerNormalization(epsilon=1e-6) # final norm layer
self.head = Dense(num_classes)
else:
self.avg = None
self.norm = None
self.head = None
self.flatten = Flatten()
self.last = Dense(1)
def forward_features(self, x):
def call(self, x):
x = tf.transpose(x, [0,2,3,1])
for i in range(4):
x = self.downsample_layers[i](x)
x = self.stages[i](x)
x = self.flatten(x)
x = self.last(x)
return x
def call(self, x):
x = self.forward_features(x)
if self.include_top:
x = self.avg(x)
x = self.norm(x)
x = self.head(x)
return x
model_urls = {
convunext.py
+487
View File
@@ -0,0 +1,487 @@
import os
import time
import argparse
import wandb
import random
import numpy as np
import tensorflow as tf
import scipy.interpolate as spi
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='convunext_7'
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name=name,
notes='basic pyramid. more balanced levels. take ssim (not ms-ssim) loss at each level. higher lr. on on jumpsuits. \
uv maps filled with nearest enighbour, up/downscaling with lanczos. only take loss at final level')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder.py')
artifact.add_file('stylegan2_ae/layers/mapping_block.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
# def plot_3d(mesh_gt, mesh_pred, faces):
# mesh_pred[:,1] += 2
# fig = go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_gt[:,0],
# y=mesh_gt[:,1],
# z=mesh_gt[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_pred[:,0],
# y=mesh_pred[:,1],
# z=mesh_pred[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2],
# )
# ])
# fig.update_layout(scene_aspectmode='manual',
# scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(
# scene = dict(
# xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,3],),
# zaxis = dict(nticks=4, range=[-1,1],),)
# )
# return fig
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
def plot_3d_levels(mesh_levels, faces):
levels = len(mesh_levels)
plot_levels = mesh_levels.copy() #[mesh_levels[i][0] for i in range(levels)]
for i in range(levels):
plot_levels[i][:,1] = plot_levels[i][:,1] + i * 2
meshes = [go.Mesh3d(
# 8 vertices of a cube
x=mesh[:,0],
y=mesh[:,1],
z=mesh[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
) for mesh in mesh_levels]
fig = go.Figure(data= meshes)
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=levels, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,levels*2],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x, training=training)
return x
class PyramidLevel(tf.keras.Model):
def __init__(self, depths=[3, 3, 9, 3], dims=[64, 128, 256, 256]):
super(PyramidLevel, self).__init__()
# reverse_dims = list(reversed(dims[:-1] + [16]))
# reverse_dims = dims[:-1] + [16]
reverse_dims = [16] + dims[:-1]
self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims)
self.bottleneck = Bottleneck(dims[-1])
self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims)
def call(self, inputs, training=None):
x = inputs
x, skip_features = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder([x, skip_features], training=training)
return x
def downscale(img, uv_mask=None):
# img = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
# coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
# img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
# img_data = img[img_idx[:,0], img_idx[:,1], :]
# interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
# img_2 = interpolator(coords)
# res = np.zeros((512,256,3))
# res[coords[:,0], coords[:,1], :] = img_2
# res[img_idx[:,0], img_idx[:,1], :] = img_data
results = [img]
if uv_mask is not None:
masks = [uv_mask]
for i in range(4):
# print(np.array(img.shape[:2]) / (2 ** (i + 1)))
results.append(tf.image.resize(img, size=np.array(img.shape[1:3]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.LANCZOS5, antialias=True, preserve_aspect_ratio=True))
if uv_mask is not None:
masks.append(tf.image.resize(uv_mask[:,:,:,tf.newaxis], size=np.array(uv_mask.shape[1:]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=False)[:,:,:,0])
if uv_mask is not None:
return results, masks
return results
def upscale(imgs):
results = []
for img in imgs:
results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5))
return results
class Pyramid(tf.keras.Model):
def __init__(self):
super(Pyramid, self).__init__()
self.level_1 = PyramidLevel(depths=[27], dims=[512])
self.level_2 = PyramidLevel(depths=[9, 27], dims=[256, 256])
self.level_3 = PyramidLevel(depths=[9, 9, 27], dims=[64, 128, 256])
self.level_4 = PyramidLevel(depths=[3, 9, 27, 3], dims=[32, 64, 128, 256])
self.level_5 = PyramidLevel(depths=[3, 3, 9, 27, 9], dims=[32, 64, 128, 256, 256])
self.levels = [self.level_5, self.level_4, self.level_3, self.level_2, self.level_1]
def call(self, inputs, training=None):
x, uv_mask, uv_shaped_body = inputs
x = tf.transpose(x, [0,2,3,1])
x_levels = downscale(x)
output = self.levels[-1](x_levels[-1])
output_levels = [output]
for i in reversed(range(len(x_levels)-1)):
level_out = self.levels[i](x_levels[i])
output = tf.image.resize(output, size=(output.shape[1] * 2, output.shape[2] * 2), method=tf.image.ResizeMethod.LANCZOS5) + level_out
output_levels.append(output)
x = output
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
return x, output_levels
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
# bottleneck_size = 256
# self.decoder = ConvNeXtReverseWithSkip(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
# self.bottleneck = Bottleneck(bottleneck_size)
# self.encoder = ConvNeXtWithSkip(depths=[3, 3, 9, 3], dims=[64, 128, 256, bottleneck_size])
self.level = PyramidLevel(depths=[3, 3, 3, 3], dims=[64, 128])
def call(self, inputs, training=None):
x, uv_mask, uv_shaped_body = inputs
x = tf.transpose(x, [0,2,3,1])
x = tf.image.resize(x, size=(256, 128))
# x, skip_features = self.encoder(x, training=training)
# x = self.bottleneck(x, training=training)
# x = self.decoder([x, skip_features], training=training)
x = self.level(x)
x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
return x
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred, levels_pred = model([tf.concat((batch['uv_static'], batch['uv_body_posed']), axis=1), batch['uv_mask'], batch['uv_shaped_body']])
# pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred = uv2mesh(pred)
# mesh_pred = pred
mesh_orig = batch['mesh_unposed']
# uv_unposed_levels, uv_mask_levels = downscale(tf.transpose(batch['uv_unposed'], [0,2,3,1]), batch['uv_mask'])
ms_ssim_levels = 0 #tf.reduce_mean([tf.reduce_mean(tf.image.ssim((levels_pred[i] + 1) * uv_mask_levels[len(levels_pred) - i - 1][:,:,:,tf.newaxis], (uv_unposed_levels[len(levels_pred) - i - 1] + 1) * uv_mask_levels[len(levels_pred) - i - 1][:,:,:,tf.newaxis], max_val=10)) for i in range(len(levels_pred))])
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
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=10))
fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'] * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
# loss = fft_loss_pc # + e_loss + c_loss
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
# mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric, 'ms_ssim': ms_ssim,
'fft_loss': fft_loss, 'euclid_error_reconst': metric_2,
'euclid_error_posed': euclid_posed, 'ms_ssim_levels': ms_ssim_levels }
@tf.function
def valid_step(batch):
pred, levels_pred = model([tf.concat((batch['uv_static'], batch['uv_body_posed']), axis=1), batch['uv_mask'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
# uv_unposed_levels, uv_mask_levels = downscale(tf.transpose(batch['uv_unposed'], [0,2,3,1]), batch['uv_mask'])
ms_ssim_levels = 0#tf.reduce_mean([tf.reduce_mean(tf.image.ssim(levels_pred[i] + 1, uv_unposed_levels[len(levels_pred) - i - 1] + 1, max_val=10)) for i in range(len(levels_pred))])
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=1) * batch['uv_mask'])
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=10))
fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_unposed'] * batch['uv_mask'][:,tf.newaxis,:,:], tf.zeros_like(batch['uv_unposed'])))), axis=[1,2,3]))
loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric, 'ms_ssim': ms_ssim,
'fft_loss': fft_loss, 'euclid_error_reconst': metric_2, 'ms_ssim_levels': ms_ssim_levels
}
def plot_everything(writer, batch: dict, type: str):
pred, levels = model([tf.concat((batch['uv_static'], batch['uv_body_posed']), axis=1), batch['uv_mask'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1])
img = tf.concat((orig_img, pred_img), axis=2)
levels = upscale(levels)
mesh_pred_levels = [uv2mesh(tf.transpose(levels[i], [0,3,1,2])) for i in range(len(levels))]
with writer.as_default():
tf.summary.image('images', img, step=step, max_outputs=4)
tf.summary.image('levels', tf.concat(levels, axis=2) , step=step, max_outputs=4)
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
mesh_posed_levels = [smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_levels[i]).numpy() for i in range(len(levels))]
mesh_posed_levels = [mesh_posed_levels[j][0, mask[0], :] for j in range(len(levels))]
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + 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))]
fig_levels = plot_3d_levels(mesh_posed_levels, faces[0])
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(pred.shape[0]) }
# fig = plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i])
wandb.log( { **figs, **{ type + '/levels_garments': fig_levels } } )
if __name__ == '__main__':
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
# R = smpl.run_rodrigues((np.zeros((8,72), dtype=np.float32)))
# zero_body = smpl.run(R, np.zeros((8, 11), dtype=np.float32))[0]
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,)
}
signature_static = { 'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=8,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
batch_size=8,
signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
signature=signature_static)
# model = AutoEncoder()
model = Pyramid()
# optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8)
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.001, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
folder = 'models/ae_logs/'
train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
valid_iterator = iter(dataset_valid)
step = 0
for batch in dataset_train:
# batch = batch_to_dict(batch, signature_dynamic)
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
train_loss = train_step(batch)
for key in train_loss.keys():
with train_summary_writer.as_default():
tf.summary.scalar(key, train_loss[key], step=step)
# data_valid = 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)
for key in valid_loss.keys():
with valid_summary_writer.as_default():
tf.summary.scalar(key, valid_loss[key], step=step)
if step % 500 == 0:
plot_everything(train_summary_writer, batch, 'train')
plot_everything(valid_summary_writer, data_valid, 'valid')
# fig = make_subplots(rows=1, cols=5, start_cell="top-left")
# for i, l in enumerate(levels):
# fig.add_trace(ex.imshow((l[0] + 1) * 128).data[0], row=1, col=i+1)
print(step)
step += 1
custom_data/jumpsuit_20_avg_uv.npy
BIN
View File
Binary file not shown.
custom_utils/__init__.py → custom_utils/a__init__.py
View File
custom_utils/custom_utils.py
+49
View File
@@ -66,6 +66,33 @@ class UV2Mesh():
return self.uv_to_mesh(uv)
class UV2MeshLabel():
def __init__(self, res=256):
super(UV2MeshLabel, self).__init__()
self.res = res
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
m2uv = m2uv_array(m2uv)
self.uv_pixels = self.uv_map_discreet[m2uv]
# input = channels_first
# @tf.function
def uv_to_mesh(self, uv):
batch_size = uv.shape[0]
x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1])
x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1])
x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
mesh = (x1 + x2) / 2
return mesh
def __call__(self, uv):
return self.uv_to_mesh(uv)
def _maskFaces(V, F, mask):
V_mask = mask
@@ -79,6 +106,21 @@ def _maskFaces(V, F, mask):
return V_new, F_new
# def _maskFaces_tf(V, F, mask):
# V_mask = mask
# F_mask = tf.gather(V_mask, F.reshape(-1), axis=0)
# F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
# F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3
# V_ind = tf.zeros(V_mask.shape)
# V_ind[V_mask] = tf.range(V_mask.sum())
# F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
# V_new = V[V_mask, :]
# return V_new, F_new
def maskFacesBatch(V, F, mask):
Vs = []
Fs = []
@@ -148,3 +190,10 @@ conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
def nest_conditions(batch):
batch['conditions'] = {d: batch.pop(d) for d in conditions}
return batch
def sparse_to_tensor(L):
idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
data = L.data.astype(dtype=np.float32)
shape = L.shape
return tf.SparseTensor(idx, data, shape)
custom_utils/graph_utils.py
+88 -5
View File
@@ -1,10 +1,14 @@
import os
import sys
import scipy as sp
import numpy as np
import scipy.io as sio
import tensorflow as tf
from tensorflow.keras.activations import tanh, relu
from tensorflow.keras.layers import Dense, BatchNormalization, Add, Flatten
from custom_utils.custom_utils import UV2Mesh
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
# from custom_utils.custom_utils import UV2Mesh
from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
@@ -31,6 +35,14 @@ def repeat_laplacian(L, times):
return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
def convert_pooling(pooling, batchSize):
pool_map = np.zeros((np.max(pooling)+1, pooling.shape[0]), dtype=np.float32)
for i in range(pooling.shape[0]):
pool_map[pooling[i], i] = 1.0
pool_map = np.tile(pool_map[np.newaxis, :, :], [batchSize, 1, 1])
pool_map = tf.sparse.from_dense(pool_map)
return pool_map
class GraphConvolution(tf.keras.layers.Layer):
def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
@@ -174,8 +186,76 @@ class GNNDiscriminator(tf.keras.Model):
return x
class Reconstructor_V4_Decoder(tf.keras.Model):
def __init__(self, batchSize, isSkirt=False, n_feature=64):
super(Reconstructor_V4_Decoder, self).__init__()
laplacians = [sio.loadmat("custom_utils/laplacian/ssmpl" + ("_skirt" if isSkirt else "") + str(i) + ".mat")['L'].tocoo() for i in range(5)]
poolings = [sio.loadmat("custom_utils/pooling/pool" + str(i) + str(i+1) + ".mat")['idx'] for i in range(5)]
poolings = [convert_pooling(poolings[i][0,:], batchSize) for i in range(len(poolings))]
self.custom_layers = []
for i in range(4,1,-1):
self.custom_layers.append(GraphUnPool(pool_map=poolings[i]))
self.custom_layers.append(GraphConvolution((int(n_feature/2**(4-(i if i == 4 else i+1))), int(n_feature/2**(4-i)) ), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
for _ in range(2**(i-2)):
conv1 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
conv2 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
self.custom_layers.append([conv1, conv2])
self.pool2 = GraphUnPool(pool_map=poolings[1])
self.conv2 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**2)), activation=relu, L=laplacians[1], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_2')
self.pool1 = GraphUnPool(pool_map=poolings[0])
self.conv1 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**3)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
self.conv0 = GraphConvolution((int(n_feature/2**3), 3), L=laplacians[0], batchSize=batchSize, name='gcn_0')
def call(self, inputs, training=None):
x = inputs
for layer in self.custom_layers:
if isinstance(layer, list):
x1 = layer[0](x, training)
x2 = layer[1](x1, training)
x = Add()([x, x2])
else:
x = layer(x, training=training)
x = self.pool2(x)
x = self.conv2(x, training)
x = self.pool1(x)
x = self.conv1(x, training)
x = self.conv0(x, training)
return x
class Reconstructor_v2(tf.keras.Model):
def __init__(self, batchSize, isSkirt=False):
super(Reconstructor_v2, self).__init__()
self.flatten = tf.keras.layers.Flatten()
self.fc2 = tf.keras.layers.Dense(14475*3, kernel_regularizer=tf.keras.regularizers.l2)
self.reshape2 = tf.keras.layers.Reshape((14475, 3))
# self.drop = Dropout(0.2)
# self.fc3 = Dense(512)
self.fc1 = tf.keras.layers.Dense(3)
def call(self, inputs, training=None):
x = self.flatten(inputs)
x = self.fc2(x)
# x = self.drop(x, training=training)
x = self.reshape2(x)
# x = self.fc3(x)
x = self.fc1(x)
return x
class DatasetTest(tf.data.Dataset):
@@ -195,13 +275,16 @@ class DatasetTest(tf.data.Dataset):
if __name__ == '__main__':
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
batch_size = 4
input = tf.ones((batch_size, 14475, 3))
# input = tf.ones((batch_size, 14475, 3))
ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
# ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
x = ge(input)
# x = ge(input)
gd = Reconstructor_V4_Decoder(batch_size)
x = gd(tf.ones((batch_size, 21, 64)))
print()
custom_utils/laplacian/laplacian0.npz
BIN
View File
Binary file not shown.
custom_utils/laplacian/laplacian5.npz
BIN
View File
Binary file not shown.
custom_utils/smpl_faces_skirt.npy
BIN
View File
Binary file not shown.
cvae.py
+584
View File
@@ -0,0 +1,584 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# Currently, memory growth needs to be the same across GPUs
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
except RuntimeError as e:
# Memory growth must be set before GPUs have been initialized
print(e)
import time
import argparse
import wandb
import random
import numpy as np
import scipy as sp
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
from auto_encoder_body import AutoEncoder as BodyAE
from auto_encoder_curv import AutoEncoder as CurvatureAE
from auto_encoder_static import AutoEncoder as StaticAE
name='cvae_10_2'
def init_wandb():
run = wandb.init(project='vae', name=name,
notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline.\
new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1000 from 250.\
offsets over average unposed garment. gamma balancing. remove global orientation (in posed body input and in pose_shape parameters. larger batchsize\
add edge and collision (fixed) loss. rebalanced l1 loss. increasing edge and coll coeff. retrained body encoder with deoriented bodies \
don\'t freeze encoders, lower lr')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x, training=training)
return x
class Encoder(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder, self).__init__()
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder(tf.keras.Model):
def __init__(self):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * 80)
self.reshape = tf.keras.layers.Reshape((16, 8, 80))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class CVAE(tf.keras.Model):
def __init__(self, latent_dim=128):
super(CVAE, self).__init__()
self.latent_dim = latent_dim
model : StaticAE = StaticAE()
load_model(model, 'checkpoints/ae_static_3')
self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
# self.static_enc.trainable = False
model : BodyAE = BodyAE()
load_model(model, 'checpoints/ae_body_5')
self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
# self.body_enc.trainable = False
model : CurvatureAE = CurvatureAE()
load_model(model, 'checpoints/ae_curvature_6')
self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
# self.curv_enc.trainable = False
self.encoder = Encoder(latent_dim)
self.decoder = Decoder()
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, eps=None):
if eps is None:
eps = tf.random.normal(shape=(4, self.latent_dim))
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [4,1,1,1])
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [4,1,1])
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = x + avg_jumpsuit[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def decode(self, z, apply_sigmoid=False):
logits = self.decoder(z)
if apply_sigmoid:
probs = tf.sigmoid(logits)
return probs
return logits
def call(self, inputs, training=None):
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
z_static = self.static_enc(uv_template)
z_body = self.body_enc(uv_body_posed)
z_curv = self.curv_enc(uv_curv)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z = tf.concat([z, z_static, z_body, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
x = x + avg_jumpsuit[tf.newaxis,:,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
return res
def kl_loss(mean, logvar):
kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
kl_loss = tf.reduce_mean(kl_loss)
kl_loss *= -0.5
return kl_loss
@tf.function
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'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
mesh_orig = batch['mesh_unposed']
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
# 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))
rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
tv_loss = tf.image.total_variation(pred)
reconstr_loss = l1_loss / 4
reconstr_loss = tf.reduce_mean(reconstr_loss) + reg_coeff *(e_loss + c_loss)
# logpx_z = -reconstr_loss
# logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
# logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
# loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
kl = kl_loss(mean, logvar)
loss = reconstr_loss / (2 * (gamma ** 2)) + balancing_coeff * kl
# loss = reconstr_loss + balancing_coeff * kl
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
fft_loss = tf.reduce_mean(fft_loss)
tv_loss = tf.reduce_mean(tv_loss)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
'train/reg_coeff': reg_coeff }
@tf.function
def valid_step(batch):
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'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_orig = batch['mesh_unposed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
# mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
reconstr_loss = l1_loss
# logpx_z = -reconstr_loss
# logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
# logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
# loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
kl = kl_loss(mean, logvar)
# loss = reconstr_loss / (2 * (gamma ** 2)) + kl
loss = reconstr_loss + balancing_coeff * kl
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
fft_loss = tf.reduce_mean(fft_loss)
b_loss = tf.reduce_mean(b_loss)
reconstr_loss = tf.reduce_mean(reconstr_loss)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
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'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
pred_img = pred
img = tf.concat((uv_static_img, uv_body_posed_img, 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")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + 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[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def load_model(model, checkpoint):
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
ckpt.restore(manager.latest_checkpoint)
def sample(folder):
load_model(model,folder)
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0])
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [4,1]), tf.tile(shape[tf.newaxis,:], [4,1])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
mesh_orig = batch['mesh_posed'][0].numpy()
body = batch['mesh_body_posed'][0].numpy()
body_color = np.ones_like(body, dtype=np.int32)*100
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
batch = 0
meshes = []
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=mesh_orig[mask,0],
y=mesh_orig[mask,1],
z=mesh_orig[mask,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(4):
x = mesh_posed[i][mask, :].copy()
x[:, 1] += 1 * (i + 1)
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(5):
x = body.copy()
x[:, 1] += 1 * i
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = F[:,0],
j = F[:,1],
k = F[:,2],
vertexcolor=body_color
))
fig = go.Figure(data=meshes)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=3, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,5],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
fig.show()
print()
if __name__ == '__main__':
# init_wandb()
# allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
batch_size = 12
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
'uv_curvature': (512, 256),
}
signature_static = { 'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model = CVAE()
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=0.0001)
balancing_coeff = 250
reg_coeff = 1
folder = 'models/ae_logs/'
valid_iterator = iter(dataset_valid)
gamma = 1
balancing_activated_step = 0
avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
sample('checkpoints/cvae_10')
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
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)
valid_loss = {}
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 2500 == 0:
train_figs = plot_everything(batch, 'train')
# valid_figs = plot_everything(data_valid, 'valid')
valid_figs = {}
to_log = { **to_log, **train_figs, **valid_figs }
if step % 10000 == 0:
manager.save(checkpoint_number=step)
if train_loss['train/euclid_error'] < 0.01:
if balancing_activated_step == 0:
balancing_activated_step = step
balancing_coeff = balancing_coeff + 1000 / 30000
if train_loss['train/euclid_error'] < 0.015:
reg_coeff = reg_coeff + 100 / 15000
# if step > 15000 and balancing_coeff <= 1e-3:
# balancing_coeff = 1e-2
# if step > 20000 and balancing_coeff <= 1e-3:
# balancing_coeff = 1e-2
# if step > 25000 and balancing_coeff <= 1e-2:
# balancing_coeff = 1e-1
# if step > 30000 and balancing_coeff <= 1e-1:
# balancing_coeff = 1
# gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
wandb.log(to_log, step = step)
print(step)
step += 1
dataset.py
+231 -39
View File
@@ -6,7 +6,7 @@ import random
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import nest_conditions
from custom_utils.custom_utils import nest_conditions, batch_to_dict
@@ -101,7 +101,7 @@ class DatasetIdx(tf.data.Dataset):
def __new__(cls, folder):
uv_static_folder = os.path.join(folder, 'pose_shape')
frames = os.listdir(uv_static_folder)
frames = sorted(os.listdir(uv_static_folder))#[453:503]
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
random.shuffle(frames)
@@ -119,6 +119,8 @@ class DatasetFactory():
for key in signature.keys():
if key == 'id':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'uv_curvature' or key == 'mesh_curvature':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
else:
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
return tf_signature
@@ -139,16 +141,27 @@ class DatasetFactory():
item['id'] = id
if 'pose_shape' in signature.keys():
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
##############################
pose_shape[:3] = 0
##############################
item['pose_shape'] = pose_shape
if 'uv_posed' in signature.keys():
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
item['uv_posed'] = uv_posed
# if 'uv_body_posed' in signature.keys():
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
# item['uv_body_posed'] = uv_body_posed
# if 'mesh_body_posed' in signature.keys():
# mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
# item['mesh_body_posed'] = mesh_body_posed
if 'uv_body_posed' in signature.keys():
uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
item['uv_body_posed'] = uv_body_posed
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
item['uv_body_posed'] = uv_posed_body_no_orient
if 'mesh_body_posed' in signature.keys():
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
item['mesh_body_posed'] = mesh_body_posed
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
item['mesh_body_posed'] = mesh_posed_body_no_orient
if 'mesh_posed' in signature.keys():
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
item['mesh_posed'] = mesh_posed
@@ -158,7 +171,12 @@ class DatasetFactory():
if 'uv_unposed' in signature.keys():
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
item['uv_unposed'] = uv_unposed
if 'uv_curvature' in signature.keys():
uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature'] = uv_curvature.astype(np.int32)
if 'mesh_curvature' in signature.keys():
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
# item = { **static_info, **item }
as_list = []
@@ -168,6 +186,10 @@ class DatasetFactory():
return as_list
# dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
# inp=[i],
# Tout=self._build_tf_signature(signature)
# ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i],
Tout=self._build_tf_signature(signature)
@@ -186,44 +208,232 @@ class DatasetFactory():
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
if __name__ == '__main__':
from custom_utils.custom_utils import batch_to_dict
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
import plotly.graph_objects as go
# import plotly.express as ex
# import cmapy as cm
uv2mesh = UV2Mesh()
F = np.load("custom_utils/smpl_faces.npy")
# dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
# for d in dataset:
# print(d)
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
# 'uv_curvature': (512, 256),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
# 'uv_mask': (512, 256,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train',
batch_size=1024,
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=1,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train',
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
@tf.function
# sum_uv = np.zeros((3, 512, 256), dtype=np.float32)
# sum_mask = np.zeros((512, 256), dtype=np.float32)
# for i, batch in enumerate(dataset):
# print(i)
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
# sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
# sum_mask += batch['uv_mask'].numpy().sum(axis=0)
# avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :]
# sum_mask[sum_mask == 0] = 1
# x = np.transpose(avg_jumpsuit, [1,2,0])
# x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
# go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,1],
# z=x[:,2],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2]
# )]).show()
# np.save('custom_data/jumpsuit_20_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
# print()
uv_res = 256
def m2uv_array(m2uv):
result = np.zeros((14475, 2), dtype=np.int32)
for k, v in m2uv.items():
if len(v) == 2:
result[k] = list(v)
elif len(v) == 1:
x = v.pop()
result[k] = [x, x]
else:
assert False
return result
def unique_uv_map_mask(uv_map_discreet):
mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
temp = np.zeros((uv_res * 2, uv_res))
for p in uv_map_discreet:
temp[p[0], p[1]] += 1
# temp = temp == 1
for i, p in enumerate(uv_map_discreet):
if temp[p[0], p[1]] == 1:
mask[i] = True
elif temp[p[0], p[1]] == 2:
mask[i] = True
temp[p[0], p[1]] = 0
return mask
def uv_to_mesh(uv, m2uv):
res = uv.shape[0]
mesh = np.zeros((14475, 3))
uv_pixels = uv_map_discreet[m2uv]
# counter = np.ones((14475))
# counter[m2uv[:, 0] == m2uv[:, 1]] = 2
mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
return mesh
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
m2uv = m2uv_array(m2uv)
uv_mask = np.load('custom_utils/mask_256x512.npy')
curvature_label_type = '6_25'
uv_res = 256
uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
uv_map_mask = unique_uv_map_mask(uv_map_discreet)
uv_counter = np.zeros((uv_map.shape[0]))
uv_counter[m2uv[:, 0]] = 1
uv_counter[m2uv[:, 1]] += 1
uv_counter = uv_counter[uv_map_mask]
for i, batch in enumerate(dataset):
print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
values = mesh_unposed
uv_values = np.zeros((uv_map.shape[0], 3))
uv_values[m2uv[:, 0]] = values
uv_values[m2uv[:, 1]] += values
uv_counter = np.zeros((uv_map.shape[0]))
uv_counter[m2uv[:, 0]] = 1
uv_counter[m2uv[:, 1]] += 1
uv_values /= uv_counter[:, np.newaxis]
x = uv_values
# verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
im = ((x - x.min()) / (x.max() - x.min())) * 255
x = np.transpose(uv_unposed, [1,2,0])
x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
)]).show()
print()
# @tf.function
def euclid_error(mask, gt, pred):
return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))) / tf.reduce_sum(mask)
# return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1)
return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
all_error = []
for batch in dataset:
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], batch['mesh_body_posed'])
all_error.append(err.numpy())
mesh_reconstruct = uv2mesh(batch['uv_posed'])
err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
i = 296
err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
_, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=y[:,0],
y=y[:,1],
z=y[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
)])
fig.show()
fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
fig.update_traces(marker={'color': colors, 'size': 5})
fig.show()
all_error.append(err.numpy().sum())
# b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
print(len(all_error))
@@ -232,25 +442,7 @@ if __name__ == '__main__':
from losses import collision_loss as collision_penalty
from losses import collision_penalty as collision_penalty_2
from custom_utils.custom_utils import UV2Mesh
import matplotlib.pyplot as plt
def plot_3d(x, y=None):
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(111, projection='3d')
ax.set_xlim3d(-1, 1)
ax.set_ylim3d(-1, 1)
ax.set_zlim3d(-1, 1)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.grid(False)
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
if y is not None:
ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
plt.savefig('temp.jpg')
uv2mesh = UV2Mesh()
for d in dataset:
d
docker/Dockerfile
+1 -1
View File
@@ -3,7 +3,7 @@ FROM tensorflow/tensorflow:2.9.0-gpu
RUN apt-get update
RUN apt-get upgrade -y
RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python
RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python cmapy
WORKDIR /home/workspace/data
docker/requierements.txt
+2 -1
View File
@@ -8,4 +8,5 @@ scipy
plotly
pandas
tensorboardX
opencv-python
opencv-python
cmapy
downscaling_test.py
+259
View File
@@ -0,0 +1,259 @@
import os
import time
import argparse
import wandb
import random
import tqdm
import scipy as sp
import scipy.interpolate as spi
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
uv2mesh = UV2Mesh(res=256)
def downscale_1(img):
results = [img]
for i in range(5):
results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.LANCZOS5))
return results
def downscale_2(img):
results = [img]
for i in range(5):
# print(np.array(img.shape[:2]) / (2 ** (i + 1)))
results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
return results
def upscale_1(imgs):
results = []
for img in imgs:
results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5))
return results
def optimize(img, gt):
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1.0, decay_steps=40000, decay_rate=0.1, staircase=False)
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
img = tf.Variable(img)
@tf.function
def error():
pred = tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
loss = 1 - tf.image.ssim_multiscale(pred + 1, gt + 1, max_val=2)
return loss
for i in range(1000):
optimizer.minimize(error, [img])
if i % 100 == 0:
pred = tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
loss = 1 - tf.image.ssim_multiscale(pred + 1, gt + 1, max_val=2)
print('it: ' + str(i) + ' loss: ' + str(loss.numpy()))
return img
def fill(img, mask):
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
img_data = img[img_idx[:,0], img_idx[:,1], :]
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
img_2 = interpolator(coords)
res = np.zeros((512,256,3), dtype=np.float32)
res[coords[:,0], coords[:,1], :] = img_2
res[img_idx[:,0], img_idx[:,1], :] = img_data
return res
def fill_uv_maps(folder, type):
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
mask = np.load('custom_utils/mask_256x512.npy')
frames = os.listdir(os.path.join(folder, type))
for frame in tqdm.tqdm(frames):
id, _ = frame.split('_')
# mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
img = np.load(os.path.join(folder, type, frame)); img = img['uv_body']
filled = fill(img, mask)
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_body_posed = filled)
if __name__ == '__main__':
# fill_uv_maps(folder='/home/workspace/data/dataset/stylegan_jumpsuit/train', type='uv_body_posed')
# pass
# exit(0)
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train',
batch_size=1,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
signature=signature_static)
step = 0
for batch in dataset_train:
# batch = batch_to_dict(batch, signature_dynamic)
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
img = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_data = img[img_idx[:,0], img_idx[:,1], :]
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
img_2 = interpolator(coords)
res = np.zeros((512,256,3))
res[coords[:,0], coords[:,1], :] = img_2
res[img_idx[:,0], img_idx[:,1], :] = img_data
img_2 = res
down_1 = downscale_1(img)
down_2 = downscale_1(img_2)
up_1 = upscale_1(down_1)
up_2 = upscale_1(down_2)
mask = batch['mesh_mask'].numpy().astype(bool)[0]
mesh_2 = [uv2mesh(tf.transpose(up_2[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_2))]
mesh_1 = [uv2mesh(tf.transpose(up_1[i], [2,0,1])[tf.newaxis,:,:,:])[0] for i in range(len(up_1))]
F = np.load("custom_utils/smpl_faces.npy")
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
# mesh_downscaled = mesh_2[-1].numpy()[mask, :]
# uv_optimized = optimize(down_2[-1], tf.transpose(batch['uv_posed'][0], [1,2,0]))
# uv_optimized = tf.image.resize(uv_optimized, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
# mesh_optimized = uv2mesh(tf.transpose(uv_optimized, [2,0,1])[tf.newaxis,:,:,:])[0].numpy()[mask, :]
# go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_downscaled[:,0],
# y=mesh_downscaled[:,1],
# z=mesh_downscaled[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_optimized[:,0],
# y=mesh_optimized[:,1],
# z=mesh_optimized[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2],
# )
# ]).show()
# fig = make_subplots(rows=2, cols=6, start_cell="top-left")
# for i, d1 in enumerate(down_1):
# fig.add_trace(ex.imshow((d1 + 1) * 128).data[0], row=1, col=i+1)
# for i, u1 in enumerate(up_1):
# fig.add_trace(ex.imshow((u1 + 1) * 128).data[0], row=2, col=i+1)
# fig.show()
fig = make_subplots(rows=2, cols=6, start_cell="top-left")
for i, d2 in enumerate(down_2):
fig.add_trace(ex.imshow((d2 + 1) * 128).data[0], row=1, col=i+1)
for i, u2 in enumerate(up_2):
fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=2, col=i+1)
fig.show()
fig = make_subplots(rows=2, cols=6, start_cell="top-left", specs = [ [{ 'type': 'xy' } for _ in range(6)], [{ 'type': 'scene' } for _ in range(6)] ])
for i, u2 in enumerate(up_2):
fig.add_trace(ex.imshow((u2 + 1) * 128).data[0], row=1, col=i+1)
# for i, m1 in enumerate(mesh_1):
# m1 = m1.numpy()[mask, :]
# fig.add_trace(go.Mesh3d(
# # 8 vertices of a cube
# x=m1[:,0],
# y=m1[:,1],
# z=m1[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
# row=1, col=i+1)
for i, m2 in enumerate(mesh_2):
m2 = m2.numpy()[mask, :]
fig.add_trace(go.Mesh3d(
# 8 vertices of a cube
x=m2[:,0],
y=m2[:,1],
z=m2[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
row=2, col=i+1)
fig.show()
print()
downscaling_test_curv.py
+227
View File
@@ -0,0 +1,227 @@
import os
import time
import argparse
import wandb
import random
import tqdm
import scipy as sp
import cmapy as cm
import scipy.interpolate as spi
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2MeshLabel, maskFacesBatch, _maskFaces
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
uv2mesh = UV2MeshLabel(res=256)
def downscale_1(img):
results = [img]
for i in range(5):
results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.LANCZOS5))
return results
def downscale_2(img):
results = [img]
for i in range(5):
# print(np.array(img.shape[:2]) / (2 ** (i + 1)))
results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True))
return results
def upscale_1(imgs):
results = []
for img in imgs:
results.append(tf.image.resize(img, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5))
return results
def fill(img, mask):
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
img_data = img[img_idx[:,0], img_idx[:,1], :]
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
img_2 = interpolator(coords)
res = np.zeros((512,256,3), dtype=np.float32)
res[coords[:,0], coords[:,1], :] = img_2
res[img_idx[:,0], img_idx[:,1], :] = img_data
return res
def fill_uv_maps(folder, type):
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
mask = np.load('custom_utils/mask_256x512.npy')
frames = os.listdir(os.path.join(folder, type))
for frame in tqdm.tqdm(frames):
id, _ = frame.split('_')
# mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
img = np.load(os.path.join(folder, type, frame)); img = img['uv_body']
filled = fill(img, mask)
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_body_posed = filled)
if __name__ == '__main__':
# fill_uv_maps(folder='/home/workspace/data/dataset/stylegan_jumpsuit/train', type='uv_body_posed')
# pass
# exit(0)
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
'uv_curvature': (512, 256),
'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train',
batch_size=1,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
signature=signature_static)
step = 0
for batch in dataset_train:
# batch = batch_to_dict(batch, signature_dynamic)
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
img = tf.squeeze(batch['uv_curvature']).numpy()
coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_data = img[img_idx[:,0], img_idx[:,1]]
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
img_2 = interpolator(coords)
res = np.zeros((512,256))
res[coords[:,0], coords[:,1]] = img_2
res[img_idx[:,0], img_idx[:,1]] = img_data
img_2 = res
i = 0
down_curv = tf.image.resize(img_2[:,:,tf.newaxis], size=np.array(img_2.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
down_curv = tf.image.resize(down_curv, size=(512, 256), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)[:,:,0]
img_posed = tf.transpose(tf.squeeze(batch['uv_posed']), [1,2,0]).numpy()
coords = np.transpose(np.stack(np.nonzero((1 - batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_idx = np.transpose(np.stack(np.nonzero((batch['uv_mask'][0]).numpy().astype(bool))), (1,0))
img_data = img_posed[img_idx[:,0], img_idx[:,1], :]
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
img_posed_2 = interpolator(coords)
res = np.zeros((512,256,3))
res[coords[:,0], coords[:,1], :] = img_posed_2
res[img_idx[:,0], img_idx[:,1], :] = img_data
img_posed_2 = res
mesh_posed = uv2mesh(tf.transpose(img_posed_2[i], [2,0,1])[tf.newaxis,:,:,:])[0]
mask = batch['mesh_mask'].numpy().astype(bool)[0]
mesh_1 = uv2mesh(img_2[tf.newaxis,:,:])[0].numpy()[mask]
mesh_2 = uv2mesh(down_curv[tf.newaxis,:,:])[0].numpy()[mask]
F = np.load("custom_utils/smpl_faces.npy")
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
fig = make_subplots(rows=1, cols=2, start_cell="top-left")
fig.add_trace(ex.imshow((img_2) / 6).data[0], row=1, col=1)
fig.add_trace(ex.imshow((down_curv ) / 6).data[0], row=1, col=2)
fig.show()
mesh = tf.squeeze(batch['mesh_posed']).numpy()[mask]
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
v_color_1 = np.zeros((mesh.shape[0], 3))
v_color_2 = np.zeros((mesh.shape[0], 3))
for i in range(6+1):
v_color_1[mesh_1 == i, :] = colors[i]
v_color_2[mesh_2 == i, :] = colors[i]
mesh_copy = mesh.copy()
mesh_copy[:,1] += 1
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh[:,0],
y=mesh[:,1],
z=mesh[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
vertexcolor=v_color_1,
text='original'
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_copy[:,0],
y=mesh_copy[:,1],
z=mesh_copy[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
vertexcolor=v_color_2,
text='rescaled'
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
fig.show()
print()
load_models.py
+12 -7
View File
@@ -2,10 +2,13 @@ import tensorflow as tf
def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=True):
if custom_cuda:
from stylegan2.generator import Generator
else:
from stylegan2_ref.generator import Generator
# if custom_cuda:
# from stylegan2.generator import Generator
# else:
# from stylegan2_ref.generator import Generator
from stylegan2.generator import ConvNeXtReverseConditional as Generator
if g_params is None:
g_params = {
@@ -23,7 +26,7 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
test_conditions = { 'pose_shape': tf.ones((1, 83), dtype=tf.float32), 'uv_static': tf.ones((1, 3, 512, 256), dtype=tf.float32), 'uv_mask': tf.ones((1, 512, 256), dtype=tf.float32), 'uv_shaped_body': tf.ones((1, 3, 512, 256), dtype=tf.float32) }
# build generator model
generator = Generator(g_params)
generator = Generator()
_ = generator([test_latent, test_conditions])
if ckpt_dir is not None:
@@ -44,7 +47,8 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
# else:
# from stylegan2_ref.discriminator import Discriminator
from custom_utils.graph_utils import GNNDiscriminator as Discriminator
# from custom_utils.graph_utils import GNNDiscriminator as Discriminator
from convnext.models.convnext_tf import ConvNeXtDiscriminator as Discriminator
res = d_params['resolutions'][-1]
@@ -52,7 +56,8 @@ def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
# build discriminator model
discriminator = Discriminator(batch_size, 32)
# discriminator = Discriminator(batch_size, 32)
discriminator = Discriminator(depths=[3, 3, 27, 3], dims=[96, 192, 256, 8], include_top=False)
_ = discriminator(test_images)
if ckpt_dir is not None:
losses.py
+56
View File
@@ -6,6 +6,11 @@ import scipy.io as sio
from scipy.spatial import cKDTree
from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
from custom_utils.custom_utils import maskFacesBatch, _maskFaces
from custom_utils.graph_utils import sparse_to_tensor
def d_logistic(real_images, generator, discriminator, z_dim, conditions):
batch_size = tf.shape(real_images)[0]
@@ -202,6 +207,57 @@ def uv_mask_loss(pred, uv_mask):
return invalid_count / total_count
def merge_verts(Vs, batchDim):
V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
return V
def _bend_loss_single(x, F, L, mask):
verts = tf.boolean_mask(x, mask, axis=0)
VN = vertex_normals(verts, F)
bend = edge_convolution_template(
VN, L, sizes=None,
edge_function=lambda x,y: (x - y)**2,
reduction='weighted',
edge_function_kwargs={}
)
loss = tf.reduce_sum(bend)
return loss
def bend_loss(x, F, L, mask):
# x: predicted outfits
# F: faces
# L: laplacian
loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])])
return loss
# def bend_loss(x, F, L, mask):
# # x: predicted outfits
# # F: faces
# # L: laplacian
# # x = x * mask[:,:,tf.newaxis]
# batch_size = x.shape[0]
# verts, faces = maskFacesBatch(x, F, mask)
# nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)]
# x = tf.concat(verts, axis=0)
# F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)])
# F = tf.concat((faces[0], F), axis=0)
# VN = vertex_normals(x, F)
# bend = edge_convolution_template(
# VN, L, sizes=None,
# edge_function=lambda x,y: (x - y)**2,
# reduction='weighted',
# edge_function_kwargs={}
# )
# loss = tf.reduce_sum(bend, axis=[1])
# loss = tf.reshape(loss, (batch_size, 14475))
# loss = tf.reduce_mean(loss, axis=1)
# return loss
if __name__ == '__main__':
V = tf.ones((4, 14475, 3))
B = tf.zeros((4, 14475, 3))
stylegan2/generator.py
+92
View File
@@ -1,10 +1,19 @@
import tensorflow as tf
import numpy as np
from stylegan2.utils import lerp
from stylegan2.layers.mapping_block import Mapping
from stylegan2.layers.synthesis_block import Synthesis
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
GlobalAveragePooling2D, Layer,
LayerNormalization, Flatten, Reshape)
# from stylegan2.layers.mapping_block import Mapping
from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding
from convnext.models.convnext_tf import Block, DropPath
class Generator(tf.keras.Model):
def __init__(self, g_params, **kwargs):
super(Generator, self).__init__(**kwargs)
@@ -115,3 +124,86 @@ class Generator(tf.keras.Model):
def compute_output_shape(self, input_shape):
return input_shape[0][0], 3, self.resolutions[-1], self.resolutions[-1]
class ConvNeXtReverseConditional(tf.keras.Model):
r""" ConvNeXt
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
https://arxiv.org/pdf/2201.03545.pdf
Args:
num_classes (int): Number of classes for classification head. Default: 1000
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
include_top (bool): whether to add head or just use it as feature extractor. Default: True
drop_path_rate (float): Stochastic depth rate. Default: 0.
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
"""
def __init__(self,
depths=[3, 3, 9, 3], dims=[64, 64, 32, 4],
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
):
super().__init__()
self.pose_shape_embedding = PoseShapeEmbedding()
self.garment_embedding = GarmentEmbedding((128 * 3) - (216 + 11))
self.reshape = Reshape((16, 8, 4))
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
stem = tf.keras.Sequential([
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
LayerNormalization(epsilon=1e-6)]
)
self.upsample_layers.append(stem)
for i in range(3):
downsample_layer = tf.keras.Sequential([
LayerNormalization(epsilon=1e-6),
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
)
self.upsample_layers.append(downsample_layer)
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(4):
stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
)
self.stages.append(stage)
cur += depths[i]
self.last = Dense(3, activation='tanh')
def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None,):
latents, conditions = inputs
pose_shape_latent = self.pose_shape_embedding([conditions['pose_shape'][:, :72], conditions['pose_shape'][:, 72:]])
garment_latent = self.garment_embedding(conditions['uv_static'])
x = tf.concat([latents, pose_shape_latent, garment_latent], axis=1)
x = self.reshape(x)
for i in reversed(range(4)):
x = self.upsample_layers[i](x)
x = self.stages[i](x)
x = self.last(x)
x = tf.transpose(x, [0,3,1,2])
x = x + conditions['uv_shaped_body']
x = x * conditions['uv_mask'][:,tf.newaxis,:,:]
return x
def set_as_moving_average_of(self, src_net):
return
def update_moving_average_of_w(self, w_broadcasted):
return
def style_mixing_regularization(self, latents1, conditions, w_broadcasted1):
return None
stylegan2/layers/embedding.py
+1 -1
View File
@@ -91,7 +91,7 @@ class GarmentEmbedding(tf.keras.layers.Layer):
# TODO: channels_first?
self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, self.embed_dim], include_top=False)
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, self.embed_dim], include_top=False)
# self.encoder(tf.keras.Input(shape=(*input_shape, 3)))
self.pooling = tf.keras.layers.GlobalAveragePooling2D()
train.py
+15 -15
View File
@@ -17,8 +17,8 @@ from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_l
from losses import collision_loss as collision_penalty
from custom_utils.smpl import SMPL
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_35', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion and edge coeff. learns offsets over shaped body. fixed unposed uv map conversion. gnn discriminator. reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input')
# os.environ["CUDA_VISIBLE_DEVICES"] = "4"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_37_2', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion and edge coeff. learns offsets over shaped body. fixed unposed uv map conversion. . reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input. convnext discriminator, convnext generator. higher starting learning rate')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
@@ -53,7 +53,7 @@ class Trainer(object):
self.train_res = t_params['train_res']
self.print_step = 100
self.save_step = 10000
self.image_summary_step = 1000
self.image_summary_step = 500
self.unsupervised_coeff = 1
self.gt_loss_coeff = 1
self.reached_max_steps = False
@@ -490,8 +490,8 @@ class Trainer(object):
test_z, test_conditions = inputs
# run networks
fake_images_05 = self.g_clone([test_z, test_conditions], truncation_psi=0.5, training=False)
fake_images_07 = self.g_clone([test_z, test_conditions], truncation_psi=0.7, training=False)
fake_images_05 = self.generator([test_z, test_conditions], truncation_psi=0.5, training=False)
fake_images_07 = self.generator([test_z, test_conditions], truncation_psi=0.7, training=False)
# merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
template = tf.concat([test_conditions['uv_static'], test_conditions['uv_static']], axis=2)
@@ -531,12 +531,12 @@ def plot_meshes(writer, garment_unposed, garment, mesh_gt, template, body, mask,
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment[i], dtype=np.int32)*230), axis=0) for i in range(len(garment))]
for i in range(len(garment)):
writer.add_mesh('garment_pc_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], global_step=step)
writer.add_mesh('garment_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
writer.add_mesh('template_' + str(i+nr_offset), vertices=template[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
writer.add_mesh('body_' + str(i+nr_offset), vertices=body[i][np.newaxis,:,:], faces=F[np.newaxis,:,:], global_step=step)
writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:])
# for i in range(len(garment)):
# writer.add_mesh('garment_pc_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], global_step=step)
# writer.add_mesh('garment_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
# writer.add_mesh('template_' + str(i+nr_offset), vertices=template[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
# writer.add_mesh('body_' + str(i+nr_offset), vertices=body[i][np.newaxis,:,:], faces=F[np.newaxis,:,:], global_step=step)
# writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:])
for i in range(len(garment)):
plotly_fig = plot_3d(mesh_gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i])
@@ -576,8 +576,8 @@ def main():
train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, args['train_res'])
_, train_d_featuremaps = filter_resolutions_featuremaps(resolutions, d_featuremaps, args['train_res'])
g_params = {
'z_dim': 256,
'w_dim': 256,
'z_dim': 128,
'w_dim': 128,
'labels_dim': -1,
'n_mapping': 8,
'resolutions': train_resolutions,
@@ -644,7 +644,7 @@ def main():
'd_params': d_params,
# training params
'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
'g_opt': {'learning_rate': 0.001, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 1e10},
'd_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16},
'batch_size': global_batch_size,
'n_total_image': 25000000,
@@ -652,7 +652,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-35-{args["train_res"]}x{args["train_res"]}')
trainer = Trainer(training_parameters, name=f'stylegan2-test-37_2-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy)
return
vae.py
+414
View File
@@ -0,0 +1,414 @@
import os
import time
import argparse
import wandb
import random
import numpy as np
import scipy as sp
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='vae_11'
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
run = wandb.init(project='vae', name=name,
notes='vae, ms_ssim with range 3. with l1 loss, larger latent_dim=1024. 4 dowscaling steps, 4 upscaling, the image resize for final resolution. upscale from 256*128')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('vae.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,1],
z=mesh_pred_unposed[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
)
def call(self, inputs, training=None):
x = inputs
x = self.bottleneck(x, training=training)
return x
class Encoder(tf.keras.Model):
def __init__(self, latent_dim):
super(Encoder, self).__init__()
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 128], include_top=False)
self.flatten = tf.keras.layers.Flatten()
self.mean = tf.keras.layers.Dense(latent_dim)
self.log_var = tf.keras.layers.Dense(latent_dim)
def call(self, inputs, training=None):
x = inputs
x = self.encoder(x)
x = self.flatten(x)
mean = self.mean(x)
log_var = self.log_var(x)
return mean, log_var
class Decoder(tf.keras.Model):
def __init__(self):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(8 * 4 * 512)
self.reshape = tf.keras.layers.Reshape((8, 4, 512))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([128, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class VAE(tf.keras.Model):
def __init__(self, latent_dim=1024):
super(VAE, self).__init__()
self.encoder = Encoder(latent_dim)
self.decoder = Decoder()
@tf.function
def sample(self, eps=None):
if eps is None:
eps = tf.random.normal(shape=(4, self.latent_dim))
return self.decode(eps, apply_sigmoid=True)
def reparameterize(self, mean, logvar):
eps = tf.random.normal(shape=mean.shape)
return eps * tf.exp(logvar * .5) + mean
def decode(self, z, apply_sigmoid=False):
logits = self.decoder(z)
if apply_sigmoid:
probs = tf.sigmoid(logits)
return probs
return logits
def call(self, inputs, training=None):
x, uv_mask = inputs
x = tf.transpose(x, [0,2,3,1])
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
return res
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_mask']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_orig = batch['mesh_unposed']
e_loss = 1e-3 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-3 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
# 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))
rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
tv_loss = tf.image.total_variation(pred)
reconstr_loss = l1_loss
logpx_z = -reconstr_loss
logpz = log_normal_pdf(z, 0., 0.)
logqz_x = log_normal_pdf(z, mean, logvar)
loss = -tf.reduce_mean(logpx_z + logpz - logqz_x)
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
fft_loss = tf.reduce_mean(fft_loss)
reconstr_loss = tf.reduce_mean(reconstr_loss)
tv_loss = tf.reduce_mean(tv_loss)
mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss }
@tf.function
def valid_step(batch):
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_mask']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_orig = batch['mesh_unposed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
# ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
# mesh_orig_2 = uv2mesh(batch['uv_unposed'])
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
reconstr_loss = l1_loss
logpx_z = -reconstr_loss
logpz = log_normal_pdf(z, 0., 0.)
logqz_x = log_normal_pdf(z, mean, logvar)
loss = -tf.reduce_mean(logpx_z + logpz - logqz_x)
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
fft_loss = tf.reduce_mean(fft_loss)
b_loss = tf.reduce_mean(b_loss)
reconstr_loss = tf.reduce_mean(reconstr_loss)
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss
}
def plot_everything(batch: dict, type: str):
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_mask']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
pred_img = pred
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")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy()
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + 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[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
return { **figs, **{ type + '/images': imgs } }
def repeat_faces(faces, times):
idx_offset = np.linspace(0, 14475 * (times-1), times).astype(np.int32)
idx_offset = np.repeat(idx_offset, faces.shape[0], axis=0)
merged = np.tile(faces, [times,1])
merged += idx_offset[:, np.newaxis]
return merged
def repeat_laplacian(L, times):
idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
idx_offset = np.linspace(0, L.shape[0] * (times-1), times).astype(np.int32)
idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
idx += idx_offset[np.newaxis,:]
data = np.repeat(L.data, repeats=times, axis=0)
shape = (L.shape[0] * times, L.shape[1] * times)
return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
if __name__ == '__main__':
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
batch_size = 8
smpl_faces = np.load('custom_utils/smpl_faces.npy')
faces_loss_single = smpl_faces
faces_loss = repeat_faces(smpl_faces, batch_size)
laplacian_loss = sp.sparse.load_npz("custom_utils/laplacian/laplacian0.npz").todense().astype(np.float32)
# edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
# laplacian_loss = sparse_to_tensor(laplacian_loss)
# laplacian_loss = sparse_to_tensor(repeat_laplacian(laplacian_loss, batch_size))
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=batch_size,
signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
batch_size=batch_size,
signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
signature=signature_static)
model = VAE()
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=lr_schedule)
folder = 'models/ae_logs/'
valid_iterator = iter(dataset_valid)
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
# _, masked_faces = maskFacesBatch(batch['mesh_unposed'].numpy(), smpl_faces, batch['mesh_mask'].numpy().astype(bool))
# # batch['mesh_masked_verts'] = masked_verts
# batch['mesh_masked_faces'] = [masked_faces[i].astype(np.int32) for i in range(batch_size)]
# batch['mesh_mask_bool'] = batch['mesh_mask'].numpy().astype(bool)
# batch['laplacian'] = [sparse_to_tensor(sp.sparse.coo_matrix(laplacian_loss[batch['mesh_mask_bool'][i], :][:, batch['mesh_mask_bool'][i]])) for i in range(batch_size)]
# batch['laplacian'] = [tf.sparse.to_dense(sparse_to_tensor(sp.sparse.coo_matrix(laplacian_loss[batch['mesh_mask_bool'][i], :][:, batch['mesh_mask_bool'][i]]))) for i in range(batch_size)]
# batch['laplacian'] = [sp.sparse.coo_matrix(laplacian_loss[batch['mesh_mask_bool'][i], :][:, batch['mesh_mask_bool'][i]]) for i in range(batch_size)]
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)
if step % 500 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **{ 'current_step': step}, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss, **{ 'current_step': step} }, step=step)
print(step)
step += 1
verify_scales.py
+81
View File
@@ -0,0 +1,81 @@
import os
import sys
import time
import cv2
import random
import numpy as np
import tensorflow as tf
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
import plotly.graph_objects as go
import plotly.express as ex
import cmapy as cm
if __name__ == '__main__':
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
uv2mesh = UV2Mesh()
F = np.load("custom_utils/smpl_faces.npy")
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
# 'uv_curvature': (512, 256),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=1,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
for i, batch in enumerate(dataset):
print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
mesh_unposed[:,1] += (1 * i)
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_unposed[mask,0],
y=mesh_unposed[mask,2],
z=mesh_unposed[mask,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
)
])
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=10, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,19],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
fig.show()
print()
verify_uv_mapping.py
+145
View File
@@ -0,0 +1,145 @@
import os
import sys
import time
from unicodedata import name
import cv2
import random
import numpy as np
import tensorflow as tf
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
import plotly.graph_objects as go
import plotly.express as ex
import cmapy as cm
if __name__ == '__main__':
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
uv2mesh = UV2Mesh()
F = np.load("custom_utils/smpl_faces.npy")
factory = DatasetFactory()
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
# 'uv_curvature': (512, 256),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=1,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
uv_res = 256
def m2uv_array(m2uv):
result = np.zeros((14475, 2), dtype=np.int32)
for k, v in m2uv.items():
if len(v) == 2:
result[k] = list(v)
elif len(v) == 1:
x = v.pop()
result[k] = [x, x]
else:
assert False
return result
def unique_uv_map_mask(uv_map_discreet):
mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
temp = np.zeros((uv_res * 2, uv_res))
for p in uv_map_discreet:
temp[p[0], p[1]] += 1
# temp = temp == 1
for i, p in enumerate(uv_map_discreet):
if temp[p[0], p[1]] == 1:
mask[i] = True
elif temp[p[0], p[1]] == 2:
mask[i] = True
temp[p[0], p[1]] = 0
return mask
def uv_to_mesh(uv, m2uv):
res = uv.shape[0]
mesh = np.zeros((14475, 3))
uv_pixels = uv_map_discreet[m2uv]
# counter = np.ones((14475))
# counter[m2uv[:, 0] == m2uv[:, 1]] = 2
mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
return mesh
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
m2uv = m2uv_array(m2uv)
uv_mask = np.load('custom_utils/mask_256x512.npy')
curvature_label_type = '6_25'
uv_res = 256
uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
uv_map_mask = unique_uv_map_mask(uv_map_discreet)
uv_counter = np.zeros((uv_map.shape[0]))
uv_counter[m2uv[:, 0]] = 1
uv_counter[m2uv[:, 1]] += 1
uv_counter = uv_counter[uv_map_mask]
for i, batch in enumerate(dataset):
print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
values = mesh_unposed
uv_values = np.zeros((uv_map.shape[0], 3))
uv_values[m2uv[:, 0]] = values
uv_values[m2uv[:, 1]] += values
uv_counter = np.zeros((uv_map.shape[0]))
uv_counter[m2uv[:, 0]] = 1
uv_counter[m2uv[:, 1]] += 1
uv_values /= uv_counter[:, np.newaxis]
x = uv_values
# verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
im = ((x - x.min()) / (x.max() - x.min())) * 255
x = np.transpose(uv_unposed, [1,2,0])
x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
)]).show()
print()