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
+26 -16
View File
@@ -36,14 +36,14 @@ from models.upscaler import UpscalerBasicConv as Upscaler
# !!!!!!!!!! CHANGE SKIRT OR NOT
low_res = 64
low_res = 128
is_skirt = True
is_cape = False
skirt_modifier = '_skirt' if is_skirt else ''
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():
run = wandb.init(project='upscale', name=name,
notes='\
@@ -112,7 +112,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
def train_step(batch):
with tf.GradientTape() as tape:
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_orig = batch['mesh_unposed']
@@ -135,7 +135,7 @@ def train_step(batch):
@tf.function
def valid_step(batch):
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_orig = batch['mesh_unposed']
@@ -156,7 +156,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
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_orig = batch['mesh_unposed']
@@ -185,9 +185,10 @@ def 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 % 100 == 0:
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)
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 10000 == 0 or (step < 2000 and step % 500 == 0):
@@ -241,15 +242,23 @@ if __name__ == '__main__':
'uv_mask' + low_resolution_modifier: (low_res, low_res//2,),
}
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)
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)
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)
# 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)
# 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_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 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)
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)
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)
if 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)
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_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:
model = Upscaler(low_res=low_res, kernel_size=3)
@@ -264,5 +273,6 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + name, max_to_keep=20)
# load_model(model, f'checkpoints/iccv/iccv_{low_res}' + skirt_modifier)
train()