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

train on template only with mask applied

Browse Source
This commit is contained in:
luni
2022-06-14 10:16:55 +00:00
parent def8a85cf9
commit 245f811d47
11 changed files with 112 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
inference_custom.py
+4 -4
View File
@@ -43,7 +43,7 @@ def uv_to_mesh(uv):
def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50):
# create generator
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
filter_index = resolutions.index(res)
g_params = {
'z_dim': 256,
@@ -66,7 +66,7 @@ def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50):
batch[1]],
training=False, truncation_psi=0.5)
np.savez_compressed(os.path.join(folder, str(i) + '.npz'), uv=tf.squeeze(fake_images), template=tf.squeeze(batch[1][2]), body=tf.squeeze(batch[3]), mask=tf.squeeze(batch[4]))
np.savez_compressed(os.path.join(folder, str(i) + '.npz'), uv=tf.squeeze(fake_images), template=tf.squeeze(batch[1][2]), body=tf.squeeze(batch[3]), mask=tf.squeeze(batch[4]), gt=tf.squeeze(batch[0]))
img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0])
cv2.imwrite(os.path.join(folder, str(i) + '.jpg'), img)
i += 1
@@ -83,9 +83,9 @@ def main():
checkpoints = [
{
'res': 512,
'ckpt_dir': 'models/stylegan2-test-8-512x512',
'ckpt_dir': 'models/stylegan2-test-12-512x512',
'use_custom_cuda': True,
'folder': 'output/test-8-196000',
'folder': 'output/test-12-140000',
'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
},
]