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

fixed full res body encoder

Browse Source
This commit is contained in:
luni
2022-09-20 22:53:02 +00:00
parent 4da4df0d61
commit d83e583e7f
1 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+12 -12
View File
@@ -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/'