import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: try: # Currently, memory growth needs to be the same across GPUs for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) logical_gpus = tf.config.experimental.list_logical_devices('GPU') print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") except RuntimeError as e: # Memory growth must be set before GPUs have been initialized print(e) import time import argparse import wandb import random import numpy as np import scipy as sp from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces import plotly.graph_objects as go import plotly.express as ex from PIL import Image from plotly.subplots import make_subplots from custom_utils.smpl import SMPL 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) 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']]) 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']]) mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) mesh_orig = batch['mesh_unposed'] orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1]) uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1]) uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1]) pred_img = pred img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] F = np.load("custom_utils/smpl_faces.npy") mask = batch['mesh_mask'].numpy().astype(bool) garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask) gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] body = batch['mesh_body_posed'].numpy() mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy() garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])] collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))] collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim) logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim) return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } def load_model(model, checkpoint): ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) ckpt.restore(manager.latest_checkpoint) 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 = [] 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: 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_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'] noise = get_noise(not random, body.shape[0], latent_dim=128) curv = model.sample(static, body, mask, noise) 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() 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) 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[:,2], z=c[:,1], 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' 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.sample(static, body, curv, uv_mask, noise=noise) mesh_pred = uv2mesh(pred) 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() pose_shape[:, :3] = 0 mesh_posed = smpl.run_garment(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped).numpy() plot_inputs(unposed[0], static[0], body[0], curv[0], mesh_mask[0]) variable_cond_data = None if random_static: variable_cond_data = static elif random_body: variable_cond_data = body elif random_curv: variable_cond_data = curv elif random_unposed: variable_cond_data = unposed else: assert False plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body) plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask) print() 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_13' folder_curv = 'checkpoints/cvae_curv_3' 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_4', body_model='checpoints/ae_body_8', curvature_model='checpoints/ae_curvature_7') curv_model = CVAECurv() load_model(model, folder_main) load_model(curv_model, folder_curv) sample(folder, variable_condition, fix_noise=True)