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

commit all

Browse Source
This commit is contained in:
luni
2023-01-25 12:25:59 +00:00
parent 1bf224841f
commit 7470ca9fb6
139 changed files with 82069 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
upscale_test.py
+17 -17
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -35,7 +35,7 @@ from models.upscaler import UpscalerBasicConv as Upscaler
# !!!!!!!!!! CHANGE SSIM BASED ON RES
# !!!!!!!!!! CHANGE SKIRT OR NOT
name='128_512_skirt'
name='256_512_skirt'
def init_wandb():
run = wandb.init(project='upscale', name=name,
notes='\
@@ -105,7 +105,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
return fig
@tf.function
# @tf.function
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)
@@ -182,7 +182,7 @@ def train():
valid_loss = valid_step(data_valid)
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
if step % 10000 == 0 or (step < 2000 and step % 500 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
@@ -204,7 +204,7 @@ def train():
step += 1
if __name__ == '__main__':
init_wandb()
# init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -214,14 +214,14 @@ if __name__ == '__main__':
# 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)
low_res = 128
low_res = 32
high_res = 512
uv_mask_body = tf.ones((low_res, low_res//2))
low_resolution_modifier = '_' + str(low_res)
high_resolution_modifier = ('_' + str(high_res)) if high_res != 512 else ''
factory = DatasetFactory()
signature_dynamic = { 'id': (),
signature_dynamic = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
'uv_unposed' + low_resolution_modifier: (low_res, low_res//2, 3),
'uv_unposed' + high_resolution_modifier: (high_res, high_res//2, 3),
@@ -234,24 +234,24 @@ if __name__ == '__main__':
'uv_mask' + low_resolution_modifier: (low_res, low_res//2,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_skirt', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_skirt', signature=signature_static, use_filled=False, 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, use_filled=False, 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, use_filled=False, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
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_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)
# dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/female', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
# static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female', signature=signature_static, use_filled=False, 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/female', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
# static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/female', signature=signature_static, use_filled=False, 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)
model = Upscaler(low_res=low_res, kernel_size=5)
model = Upscaler(low_res=low_res, kernel_size=3)
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=20)
# load_model(model, 'checkpoints/up_2_3')
load_model(model, 'checkpoints/32_512_4')
train()