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

updated body AE

Browse Source
This commit is contained in:
luni
2022-09-13 16:47:13 +00:00
parent 6ca873c76e
commit e35d7e4eb6
2 changed files with 100 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+20 -199
View File
@@ -1,4 +1,6 @@
import os import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import time import time
import argparse import argparse
import wandb import wandb
@@ -7,27 +9,19 @@ import cv2
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go import plotly.graph_objects as go
import plotly.express as ex import plotly.express as ex
from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing from losses import edge_loss, uv_mask_loss
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 dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL from custom_utils.smpl import SMPL
from models.ae_body_model import AutoEncoder
name='ae_body_7_2'
name='ae_body_8'
def init_wandb(): def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
run = wandb.init(project='vae', name=name, 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. \ 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. \ fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
@@ -36,10 +30,12 @@ def init_wandb():
larger batch_size.\ larger batch_size.\
flattened bottleneck with 128 latent_dim\ remove global orientation\ flattened bottleneck with 128 latent_dim\ remove global orientation\
fixed gender issue with deoriented bodies, again, aaand again\ fixed gender issue with deoriented bodies, again, aaand again\
trained on 100 garments') trained on 100 garments\
remove multiplications with garment uv mask')
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py') artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
artifact.add_file('models/ae_body_model.py')
artifact.add_file('dataset.py') artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py') artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact) run.log_artifact(artifact)
@@ -47,42 +43,6 @@ def init_wandb():
uv2mesh = UV2Mesh(res=256) 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): def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1 mesh_pred_unposed[:,1] += 1
@@ -133,149 +93,12 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
return fig 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 @tf.function
def train_step(batch): def train_step(batch):
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
pred = model([batch['uv_body_posed'], batch['uv_mask']]) pred = model(batch['uv_body_posed'])
# pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis] # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
@@ -283,14 +106,12 @@ def train_step(batch):
mesh_orig = batch['mesh_body_posed'] mesh_orig = batch['mesh_body_posed']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones_like(batch['mesh_mask'])))
c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], 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['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']) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed'] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1))
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)) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [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 = 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 = fft_loss_pc # + e_loss + c_loss
loss = 1e3 * (1 - ms_ssim) loss = 1e3 * (1 - ms_ssim)
@@ -304,11 +125,11 @@ def train_step(batch):
# gradients = [tf.clip_by_value(g, -1.0, 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)) 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, 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': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/euclid_error_posed': euclid_posed }
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
# UPDATE!
@tf.function @tf.function
def valid_step(batch): def valid_step(batch):
pred = model([batch['uv_body_posed'], batch['uv_mask']]) pred = model([batch['uv_body_posed'], batch['uv_mask']])
@@ -332,7 +153,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str): def plot_everything(batch: dict, type: str):
pred = model([batch['uv_body_posed'], batch['uv_mask']]) pred = model(batch['uv_body_posed'])
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed'] mesh_orig = batch['mesh_body_posed']
@@ -417,12 +238,12 @@ if __name__ == '__main__':
signature_static = { #'uv_static': (3, 512, 256,), signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,), # 'mesh_static': (14475, 3,),
'mesh_mask': (14475,), 'mesh_mask': (14475,),
'uv_mask': (512, 256,), # 'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,), # 'uv_shaped_body': (3, 512, 256,),
} }
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=False) dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=True)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic) 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) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
models/ae_body_model.py
+80
View File
@@ -0,0 +1,80 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
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