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

retrain upscalers

Browse Source
This commit is contained in:
luni
2023-02-26 12:28:31 +00:00
parent aa96c6eb9c
commit a678724711
1 changed files with 26 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
upscale_test.py
+19 -9
View File
@@ -36,14 +36,14 @@ from models.upscaler import UpscalerBasicConv as Upscaler
# !!!!!!!!!! CHANGE SKIRT OR NOT # !!!!!!!!!! CHANGE SKIRT OR NOT
low_res = 64 low_res = 128
is_skirt = True is_skirt = True
is_cape = False is_cape = False
skirt_modifier = '_skirt' if is_skirt else '' skirt_modifier = '_skirt' if is_skirt else ''
smpl_res = 3638 if is_cape else 14475 smpl_res = 3638 if is_cape else 14475
name=f'iccv_{low_res}_skirt' name=f'iccv_{low_res}_2' + skirt_modifier
def init_wandb(): def init_wandb():
run = wandb.init(project='upscale', name=name, run = wandb.init(project='upscale', name=name,
notes='\ notes='\
@@ -112,7 +112,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
def train_step(batch): def train_step(batch):
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=True) pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=True)
pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.LANCZOS5) pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.BILINEAR)
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed'] mesh_orig = batch['mesh_unposed']
@@ -135,7 +135,7 @@ def train_step(batch):
@tf.function @tf.function
def valid_step(batch): def valid_step(batch):
pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False) pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False)
pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.LANCZOS5) pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.BILINEAR)
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed'] mesh_orig = batch['mesh_unposed']
@@ -156,7 +156,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str): def plot_everything(batch: dict, type: str):
pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False) pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False)
pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.LANCZOS5) pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.BILINEAR)
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed'] mesh_orig = batch['mesh_unposed']
@@ -185,6 +185,7 @@ def train():
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
train_loss = train_step(batch) train_loss = train_step(batch)
if step % 100 == 0:
data_valid = valid_iterator.get_next() data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = valid_step(data_valid) valid_loss = valid_step(data_valid)
@@ -241,15 +242,23 @@ if __name__ == '__main__':
'uv_mask' + low_resolution_modifier: (low_res, low_res//2,), 'uv_mask' + low_resolution_modifier: (low_res, low_res//2,),
} }
if not is_skirt:
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/valid', batch_size=8, signature=signature_dynamic,resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if is_skirt:
dataset_train = factory.get_dataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) dataset_train = factory.get_dataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) static_dataset_train = StaticDataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/valid_skirt', batch_size=8, signature=signature_dynamic,resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/valid_skirt', batch_size=8, signature=signature_dynamic,resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
# dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) if is_cape:
# static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
# dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
# static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if low_res == 32 or low_res == 64: if low_res == 32 or low_res == 64:
model = Upscaler(low_res=low_res, kernel_size=3) model = Upscaler(low_res=low_res, kernel_size=3)
@@ -264,5 +273,6 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model) ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + name, max_to_keep=20) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + name, max_to_keep=20)
# load_model(model, f'checkpoints/iccv/iccv_{low_res}' + skirt_modifier)
train() train()