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

save

Browse Source
This commit is contained in:
luni
2023-03-06 09:55:48 +00:00
parent dfdf8d408e
commit 0b518e7f0e
46 changed files with 3392 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sample_pro_2.py
+35 -25
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -31,7 +31,7 @@ from PIL import Image
from plotly.subplots import make_subplots
from custom_utils.smpl import SMPL
from custom_utils.custom_utils import load_model
from custom_utils.custom_utils import load_model, save_stl
from plotly.subplots import make_subplots
from scipy.spatial.transform import Rotation
@@ -44,10 +44,12 @@ from postprocessing import fix_collision
from models.pyramid_model_optimized import Pyramid
# from configs.config_abl_overfit_dress_var_3 import config as config_main
from configs.config_abl_overfit_jumpsuit_2 import config as config_main
from configs.config_cvae_curv import config as config_curv
# from configs.config_abl_overfit_jumpsuit_2 import config as config_main
# from configs.config_cvae_curv import config as config_curv
from configs.iccv.config_abl_skin_var import config as config_main
from configs.iccv.config_cvae_curv import config as config_curv
# from models.cvae_model_overfit_jumpsuit import CVAE
from models.cvae_model import CVAE
from models.cvae_model import CVAENew as CVAE
from models.cvae_curv_model import CVAE as CVAECurv
from convert_to_uv_api import convert_to_uv_api
from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
@@ -293,17 +295,18 @@ def sampling_example(frame, camera, nr=10):
mask = batch['mesh_mask'][0].numpy().astype(bool)
_, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
noise = tf.random.normal(shape=(nr, 128))
noise = tf.random.normal(shape=(nr, 8, 4, 128))
out = model.sample(tf.tile(batch['uv_static'], [nr,1,1,1]), tf.tile(batch['uv_body_posed'], [nr,1,1,1]), tf.tile(batch['uv_curvature'], [nr,1,1,1]), tf.tile(batch['uv_mask'], [nr,1,1]), noise)
# out,_,_ = model([tf.tile(batch['uv_unposed'], [nr,1,1,1]), tf.tile(batch['uv_static'], [nr,1,1,1]), tf.tile(batch['uv_body_posed'], [nr,1,1,1]), tf.tile(batch['uv_curvature'], [nr,1,1,1]), tf.tile(batch['uv_mask'], [nr,1,1])], training=False)
posed_mesh = [get_posed_mesh(uv, batch['pose_shape'], batch['mesh_body_posed'], mask) for uv in tf.unstack(out)]
# posed_mesh = [uv2mesh(uv[tf.newaxis,:,:,:]) for uv in tf.unstack(out)]
if config_main['use_skinning']:
posed_mesh = [get_posed_mesh(uv, batch['pose_shape'], batch['mesh_body_posed'], mask) for uv in tf.unstack(out)]
else:
posed_mesh = [uv2mesh(uv[tf.newaxis,:,:,:]) for uv in tf.unstack(out)]
gt_mesh = batch['mesh_posed'][0].numpy()[mask, :]
body = batch['mesh_body_posed'][0].numpy()
posed_mesh_fixed = [postprocess(mesh, body, faces, mask, do_smoothing=False) for mesh in posed_mesh]
posed_mesh_fixed = [postprocess(mesh, body, faces, mask, do_smoothing=True) for mesh in posed_mesh]
# posed_mesh_fixed = [mesh[0].numpy()[mask, :] for mesh in posed_mesh]
# error = [np.sqrt(((gt_mesh-m) ** 2).sum(-1)).mean() for m in posed_mesh_fixed]
@@ -321,11 +324,18 @@ def sampling_example(frame, camera, nr=10):
# custom_imshow(np.concatenate(row_sampling, axis=1)).show()
os.makedirs(os.path.join(out_folder, frame), exist_ok=True)
plot_fixed(np.zeros_like(gt_mesh), body, faces, camera).write_image(os.path.join(out_folder, frame + '_body.png'))
plot_fixed(gt_mesh, body, faces, camera).write_image(os.path.join(out_folder, frame + '_gt.png'))
plot_fixed(np.zeros_like(gt_mesh), body, faces, camera).write_image(os.path.join(out_folder, frame, frame + '_body.png'))
plot_fixed(gt_mesh, body, faces, camera).write_image(os.path.join(out_folder, frame, frame + '_gt.png'))
for i in range(nr):
plot_fixed(posed_mesh_fixed[i], body, faces, camera).write_image(os.path.join(out_folder, frame + '_' + str(i) + '.png'))
plot_fixed(posed_mesh_fixed[i], body, faces, camera).write_image(os.path.join(out_folder, frame, frame + '_' + str(i) + '.png'))
save_stl(os.path.join(out_folder, frame, frame + '_body.stl'), body, F_body)
save_stl(os.path.join(out_folder, frame, frame + '_gt.stl'), gt_mesh, faces.astype(np.int32))
for i in range(nr):
save_stl(os.path.join(out_folder, frame, frame + '_' + str(i) + '.stl'), posed_mesh_fixed[i], faces.astype(np.int32))
# custom_imshow(np.concatenate(row_sampling, axis=1)).write_image(os.path.join(out_folder, frame + '_sampling.png'))
print()
@@ -418,8 +428,8 @@ if __name__ == '__main__':
model = CVAE(config_main)
curv_model = CVAECurv(config_curv)
load_model(model, 'checkpoints/' + config_main['name'])
load_model(curv_model, 'checkpoints/' + config_curv['name'])
load_model(model, 'checkpoints/iccv/' + config_main['name'])
load_model(curv_model, 'checkpoints/iccv/' + config_curv['name'])
camera_1 = dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=-0.05), eye=dict(x=0.3, y=1.0, z=0.05))
# faces_folder = '/home/workspace/data/results/support/faces'
@@ -432,12 +442,12 @@ if __name__ == '__main__':
curv_frames = ['01763_259' , '007292_282']
static_frames = ['01763_259' , '007292_282']
with open('/home/workspace/local-data/train_skirt/frames_dress_30k.txt') as f:
frames = f.readlines()
# with open('/home/workspace/local-data/train_skirt/frames_dress_30k.txt') as f:
# frames = f.readlines()
random.shuffle(frames)
frames = [frame.strip() for frame in frames]
frames = frames[:150]
# random.shuffle(frames)
# frames = [frame.strip() for frame in frames]
# frames = frames[:150]
# input_batch = get_temp_batch(frame, config_main)
# _, faces = _maskFaces(np.zeros((14475, 3)), F, input_batch['mesh_mask'][0].numpy().astype(bool))
@@ -453,11 +463,11 @@ if __name__ == '__main__':
# camera_01082_109 = dict(up=dict(x=0.0, y=0, z=1), center=dict(x=-0.15, y=0, z=-0.05), eye=dict(x=0.9, y=0.8, z=-0.1))
# frame = '04925_272'
frames = ['01070_13']
out_folder = '/home/workspace/local-data/results/sampling_plots_jumpsuits'
os.makedirs(out_folder, exist_ok=True)
for frame in tqdm.tqdm(frames):
sampling_example(frame, camera_1)
# frames = ['01070_13', '02276_232', '02978_242', '04008_295', '04173_193']
# out_folder = 'examples/sampling'
# os.makedirs(out_folder, exist_ok=True)
# for frame in tqdm.tqdm(frames):
# sampling_example(frame, camera_1)
# sampling_poses(frame, camera_01082_109)
# sampling_garment(frame, camera_1)