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()