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

pro sampling avaiable, updated cvae

Browse Source
This commit is contained in:
luni
2022-09-22 16:32:37 +00:00
parent c5ae37aacc
commit a124f95f1f
4 changed files with 297 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sample_pro.py
+218 -67
View File
@@ -1,4 +1,3 @@
from lib2to3.refactor import get_fixers_from_package
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
@@ -22,27 +21,12 @@ import random
import numpy as np
import scipy as sp
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
from tensorboardX import SummaryWriter
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from plotly.subplots import make_subplots
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
from auto_encoder_body import AutoEncoder as BodyAE
from auto_encoder_curv import AutoEncoder as CurvatureAE
from auto_encoder_static import AutoEncoder as StaticAE
from cvae_curv import CVAE as CVAECurv
from models.cvae_model import CVAE as CVAEMain
@@ -163,84 +147,251 @@ def get_frame(folder, nr):
frames = sorted(os.listdir(pose_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
random.shuffle(frames)
return frames[0], frames[1:nr]
out = [frames[0]]
out.extend(frames[1:nr])
return out
def get_bodies(folder, frame_main, frame_rand, nr=4, random=True):
def get_bodies(folder, frames, random=True):
bodies = []
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frame_main + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
pass
def get_curvs(folder, body, static, mask, frame=None, nr=4, random=True, model=None):
assert (random and model is not None) or (not random and model is None)
if random:
curv = model([body, static, mask])
uv_curvature = tf.math.argmax(pred, axis=-1)
uv_curvature = tf.one_hot(uv_curvature, 7)
pass
poses_shapes = []
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frames[0] + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
pose_shape = np.load(os.path.join(folder, 'pose_shape', frames[0] + '.npy')); poses_shapes.append(pose_shape)
if not random:
for _ in range(len(frames)-1):
bodies.append(uv_posed_body_no_orient)
poses_shapes.append(pose_shape)
else:
curv = np.load(os.path.join(folder, 'uv_curvature')) #...
pass
for i in range(1, len(frames)):
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frames[i] + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
pose_shape = np.load(os.path.join(folder, 'pose_shape', frames[i] + '.npy')); poses_shapes.append(pose_shape)
bodies = np.stack(bodies)
poses_shapes = np.stack(poses_shapes).astype(np.float32)
return bodies, poses_shapes
def get_statics(folder, nr=4, random=True):
pass
def get_unposed(folder, nr=4, random=True):
pass
def get_curvs(folder, body, static, mask, frames, noise, random=True, model: CVAECurv=None, use_original=True):
assert (random and model is not None) or (not random)
uv_curvature_orig = np.load(os.path.join(folder, 'uv_curvature', frames[0] + '.npz')); uv_curvature_orig = uv_curvature_orig['curvature_uv']
def plot_inputs(unposed, static, body, curv):
pass
noise = get_noise(not random, body.shape[0], latent_dim=128)
curv = model.sample(static, body, mask, noise)
def plot_variable_condition(condition, is_curv=False):
pass
uv_curvature = tf.math.argmax(curv, axis=-1).numpy()
if use_original:
uv_curvature[0,:,:] = uv_curvature_orig
out = tf.one_hot(uv_curvature, 7).numpy()
def plot_results(unposed, posed, body, mesh_mask):
body_color = np.ones_like(body, dtype=np.int32)*100
return out
def get_statics(folder, frames, random=True):
statics = []
uv_masks = []
mesh_masks = []
uv_static = np.load(os.path.join(folder, 'uv_static', frames[0].split('_')[0] + '.npz')); uv_static = uv_static['uv_static']; statics.append(uv_static)
uv_mask = np.load(os.path.join(folder, 'uv_mask', frames[0].split('_')[0] + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); uv_masks.append(uv_mask)
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', frames[0].split('_')[0] + '.npy')); mesh_masks.append(mesh_mask)
if not random:
for _ in range(len(frames)-1):
statics.append(uv_static)
uv_masks.append(uv_mask)
mesh_masks.append(mesh_mask)
else:
for i in range(1, len(frames)):
uv_static = np.load(os.path.join(folder, 'uv_static', frames[i].split('_')[0] + '.npz')); uv_static = uv_static['uv_static']; statics.append(uv_static)
uv_mask = np.load(os.path.join(folder, 'uv_mask', frames[i].split('_')[0] + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); uv_masks.append(uv_mask)
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', frames[i].split('_')[0] + '.npy')); mesh_masks.append(mesh_mask)
statics = np.stack(statics)
uv_masks = np.stack(uv_masks)
mesh_masks = np.stack(mesh_masks)
return statics, uv_masks, mesh_masks
def get_unposed(folder, frames, random=True):
uv_unposeds = []
mesh_unposeds = []
uv_unposed = np.load(os.path.join(folder, 'uv_posed', frames[0] + '.npz')); uv_unposed = uv_unposed['uv_gt']; uv_unposeds.append(uv_unposed)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frames[0] + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed_no_shape']; mesh_unposeds.append(mesh_unposed)
if not random:
for _ in range(len(frames)-1):
uv_unposeds.append(uv_unposed)
mesh_unposeds.append(mesh_unposed)
else:
for i in range(1, len(frames)):
uv_static = np.load(os.path.join(folder, 'uv_static', frames[i].split('_')[0] + '.npz')); uv_static = uv_static['uv_static']; uv_unposeds.append(uv_static)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frames[i] + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed_no_shape']; mesh_unposeds.append(mesh_unposed)
uv_unposeds = np.stack(uv_unposeds)
mesh_unposeds = np.stack(mesh_unposeds)
return uv_unposeds, mesh_unposeds
def make_square_2(figure):
figure.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1))
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
figure.update_layout(scene = helper,
scene2 = helper)
return figure
def make_square_4(figure):
figure.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1))
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,1],),
zaxis = dict(nticks=4, range=[-1,1],),)
figure.update_layout(scene = helper,
scene2 = helper,
scene3 = helper,
scene4 = helper,)
return figure
def plot_inputs(unposed, static, body, curv, mesh_mask):
curv = np.tile(curv.argmax(-1)[:,:,np.newaxis], [1,1,3])/6
img = np.concatenate([unposed, static, body, curv], axis=-2)
ex.imshow((img + 1) / 2).show()
mesh_static = uv2mesh(static[np.newaxis,:,:,:]).numpy()[0]
mesh_body = uv2mesh(body[np.newaxis,:,:,:]).numpy()[0]
fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
F = np.load("custom_utils/smpl_faces.npy")
_, faces = _maskFaces(np.ones((14475, 3)), F, mesh_mask)
pass
x = mesh_static[mesh_mask, :]
y = mesh_body
fig.add_trace(go.Mesh3d(
x=x[:,0],
y=x[:,1],
z=x[:,2],
i=faces[:,0],
j=faces[:,1],
k=faces[:,2]
),
row=1, col=1)
fig.add_trace(go.Mesh3d(
x=y[:,0],
y=y[:,2],
z=y[:,1],
i=F[:,0],
j=F[:,1],
k=F[:,2]
),
row=1, col=2)
make_square_2(fig)
fig.show()
def plot_variable_condition(condition, mesh_mask, is_curv=False, is_body=False):
if is_curv:
img = np.concatenate(condition.argmax(-1), axis=-1)
ex.imshow(img / img.max()).show()
else:
mesh = uv2mesh(condition).numpy()
fig = make_subplots(rows=1, cols=condition.shape[0], start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(condition.shape[0])] ])
F = np.load("custom_utils/smpl_faces.npy")
if is_body:
faces = [F for _ in range(condition.shape[0])]
else:
_, faces = maskFacesBatch(np.ones((condition.shape[0], 14475, 3)), F, mesh_mask)
for i, m in enumerate(mesh):
if not is_body:
m = m[mesh_mask[i], :]
if is_body:
x = m.copy()
m[:,1] = x[:,2]
m[:,2] = x[:,1]
fig.add_trace(go.Mesh3d(
# 8 vertices of a cube
x=m[:,0],
y=m[:,1],
z=m[:,2],
# i, j and k give the vertices of triangles
i = faces[i][:,0],
j = faces[i][:,1],
k = faces[i][:,2]
),
row=1, col=i+1)
make_square_4(fig)
fig.show()
def plot_results(unposed, posed, body, mesh_mask):
fig = make_subplots(rows=1, cols=unposed.shape[0], start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(unposed.shape[0])] ])
F = np.load("custom_utils/smpl_faces.npy")
_, faces = maskFacesBatch(np.ones((unposed.shape[0], 14475, 3)), F, mesh_mask)
body_mesh = uv2mesh(body).numpy()
collision = [np.concatenate((body_mesh[i], posed[i][mesh_mask[i], :]), axis=0) for i in range(unposed.shape[0])]
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(unposed.shape[0])]
collision_color = [np.concatenate((np.ones_like(body_mesh[i], dtype=np.int32)*100, np.ones_like(unposed[i], dtype=np.int32)*230), axis=0) for i in range(unposed.shape[0])]
for i, c in enumerate(collision):
fig.add_trace(go.Mesh3d(
x=c[:,0],
y=c[:,1],
z=c[:,2],
i=collision_faces[i][:,0],
j=collision_faces[i][:,1],
k=collision_faces[i][:,2],
vertexcolor=collision_color[i]
),
row=1, col=i+1)
make_square_4(fig)
fig.show()
def get_noise(fix_noise, nr, latent_dim=128):
eps = tf.random.normal(shape=(1 if fix_noise else nr, latent_dim))
if fix_noise:
eps = tf.tile(eps, [nr,1])
return eps
def sample(folder, variable_condition, fix_noise, nr=4):
# run = wandb.init(project='vae_plots', name=)
random_static, random_body, random_curv, random_unposed = variable_condition == 'static', variable_condition == 'body', variable_condition == 'curv', variable_condition == 'unposed'
frame_main, frame_rand = get_frame(folder, nr)
unposed, uv_mask, mesh_mask = get_unposed(folder, frame, nr, random_unposed)
static = get_statics(folder, nr, random_static)
body, pose, shape = get_bodies(folder, nr, random_body)
curv = get_curvs(folder, body, static, uv_mask, frame, nr, random_curv, curv_model)
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)
static, uv_mask, mesh_mask = get_statics(folder, frames, random_static)
unposed, mesh_unposed = get_unposed(folder, frames, random_unposed)
curv = get_curvs(folder, body, static, uv_mask, frames, noise_curv, random_curv, curv_model, use_original=True)
pred = model([unposed, static, body, curv, uv_mask], fix_noise=fix_noise)
pred = model.sample(static, body, curv, uv_mask, noise=noise)
mesh_pred = uv2mesh(pred)
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
gender = tf.cast(tf.greater_equal(pose_shape[:, -1] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(pose_shape[i,72:-1][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(len(frames))])
mesh_pred_shaped = (mesh_pred + shape_offsets).numpy()
mesh_posed = smpl.run_garment(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped).numpy()
mesh_pred_shaped = mesh_pred + shape_offsets
plot_inputs(unposed[0], static[0], body[0], curv[0], mesh_mask[0])
mesh_posed = smpl.run_garment(pose, shape, mesh_pred_shaped).numpy()
plot_inputs(unposed[0], static[0], body[0], curv[0])
variable_cond = None
variable_cond_data = None
if random_static:
variable_cond = static
variable_cond_data = static
elif random_body:
variable_cond = body
variable_cond_data = body
elif random_curv:
variable_cond = curv
variable_cond_data = curv
elif random_unposed:
variable_cond = unposed
variable_cond_data = unposed
else:
assert False
plot_variable_condition(variable_cond, is_curv=random_curv)
plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body)
plot_results(mesh_pred, mesh_posed, body, mesh_mask)
plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask)
print()
@@ -256,8 +407,8 @@ if __name__ == '__main__':
folder = '/home/workspace/local-data/train_100'
folder_main = 'checkpoints/cvae_13'
folder_curv = 'checkpoints/cvae_curv_2'
variable_condition = 'curv'
folder_curv = 'checkpoints/cvae_curv_3'
variable_condition = 'static'
batch_size = 4
smpl = SMPL()