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
sample_pro.py
+15 -95
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -28,84 +28,22 @@ from plotly.subplots import make_subplots
from custom_utils.smpl import SMPL
from custom_utils.custom_utils import load_model
from cvae_curv import CVAE as CVAECurv
from models.cvae_model import CVAE as CVAEMain
name='cvae_12_6'
def init_wandb():
run = wandb.init(project='plots', name=name,
notes='')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
from configs.config_abl_skin_var_3 import config as config_main
from configs.config_cvae_curv import config as config_curv
from models.cvae_model import CVAE
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,2] += 1
collision[:,2] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,2],
z=mesh_gt[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,2],
z=mesh_pred_unposed[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,2],
z=collision[:,1],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask'], batch['uv_shaped_body']])
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask'], batch['uv_shaped_body']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_orig = batch['mesh_unposed']
@@ -137,15 +75,6 @@ def plot_everything(batch: dict, type: str):
return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def get_frame(folder, nr):
pose_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(pose_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
random.shuffle(frames)
out = [frames[0]]
out.extend(frames[1:nr])
return out
def get_bodies(folder, frames, random=True):
bodies = []
@@ -350,11 +279,9 @@ def get_noise(fix_noise, nr, latent_dim=128):
def sample(folder, variable_condition, fix_noise, nr=4):
# run = wandb.init(project='vae_plots', name=)
def sample(folder, frames, variable_condition, fix_noise, nr=4):
random_static, random_body, random_curv, random_unposed = variable_condition == 'static', variable_condition == 'body', variable_condition == 'curv', variable_condition == 'unposed'
frames = get_frame(folder, nr)
noise = get_noise(fix_noise, nr)
noise_curv = get_noise(fix_noise, nr)
body, pose_shape = get_bodies(folder, frames, random_body)
@@ -396,32 +323,25 @@ def sample(folder, variable_condition, fix_noise, nr=4):
if __name__ == '__main__':
# init_wandb()
# allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
folder = '/home/workspace/local-data/train_100'
folder_main = 'checkpoints/cvae_18_copy'
folder_curv = 'checkpoints/cvae_curv_3'
folder = '/home/workspace/local-data/train'
variable_condition = 'static'
batch_size = 4
smpl = SMPL()
model = CVAEMain(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
latent_dim=128, static_model='checkpoints/ae_static_6_copy',
body_model='checkpoints/ae_body_8_copy',
curvature_model='checkpoints/ae_curvature_7_copy',
freeze_encoders=True)
curv_model = CVAECurv(static_model='checkpoints/ae_static_6_copy',
body_model='checkpoints/ae_body_8_copy',)
model = CVAE(config_main)
curv_model = CVAE(config_curv)
load_model(model, 'checkpoints/' + config_main['name'])
load_model(curv_model, 'checkpoints/' + config_curv['name'])
load_model(model, folder_main)
load_model(curv_model, folder_curv)
frames = ['01763_259', '007292_282', '00539_296', '00761_15']
sample(folder, variable_condition, fix_noise=True)
sample(folder, frames, variable_condition, fix_noise=True)
print()