From d83e583e7fc2238d2e2539d8add7495f4e8e837b Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 20 Sep 2022 22:53:02 +0000 Subject: [PATCH] fixed full res body encoder --- auto_encoder_body.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/auto_encoder_body.py b/auto_encoder_body.py index 0109268..775acc4 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -20,7 +20,7 @@ from custom_utils.smpl import SMPL from models.ae_body_model import AutoEncoder, AutoEncoder128, AutoEncoder32 -name='ae_body_8_d128' +name='ae_body_8' def init_wandb(): 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. \ @@ -112,7 +112,7 @@ def train_step(batch): # 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' + resolution_modifier] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1)) # 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' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2)) - ms_ssim = tf.reduce_mean(tf.image.ssim(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [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' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) 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 = fft_loss_pc # + e_loss + c_loss loss = 1e3 * (1 - ms_ssim) @@ -225,17 +225,17 @@ if __name__ == '__main__': 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) + 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) - uv_mask_body = tf.ones((128, 64)) - resolution_modifier = '_128' - resolutions = [128] + # uv_mask_body = tf.ones((128, 64)) + resolution_modifier = '' + resolutions = [512] factory = DatasetFactory() signature_dynamic = { 'id': (), 'pose_shape': (72 + 11,), - 'uv_body_posed_128': (3, 512, 256,), + 'uv_body_posed': (3, 512, 256,), 'mesh_body_posed': (14475, 3,), } signature_static = { @@ -243,14 +243,14 @@ if __name__ == '__main__': # 'uv_mask': (512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=True, resolutions=resolutions) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=resolutions) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=False, resolutions=resolutions) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=resolutions) 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() + model = AutoEncoder() # model = AutoEncoder32() - model = AutoEncoder128() + # model = AutoEncoder128() optimizer = tf.optimizers.Adam(learning_rate=0.0001) folder = 'models/ae_logs/'