import os os.environ["CUDA_VISIBLE_DEVICES"] = "4" 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 cv2 import tqdm import cmapy as cm import numpy as np import scipy as sp import trimesh as tm 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 custom_utils.custom_utils import load_model, save_stl from plotly.subplots import make_subplots from scipy.spatial.transform import Rotation from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces, load_model, group_inputs_for_levels from custom_utils.smpl import SMPL from convert_to_uv_api import convert_to_uv_api from models.upscaler import UpscalerBasicConv as Upscaler 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.iccv.config_abl_all_skirt_var import config as config_main # from configs.iccv.config_abl_all_skirt import config as config_main from configs.iccv.config_abl_all 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 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 uv2mesh = UV2Mesh(res=256, is_skirt=config_main['is_skirt']) def get_temp_batch(frame, config): with open(os.path.join(folder, 'frames_temp.txt'), 'w') as f: f.write(frame) config['signature_dynamic']['mesh_curvature'] = (14475) factory = DatasetFactory() dataset = factory.get_dataset(folder=folder, batch_size=1, signature=config['signature_dynamic'], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'], frames_file='frames_temp.txt') static_dataset = StaticDataset(folder=folder, signature=config['signature_static'], in_memory=True, is_cape=config['is_cape']) data_iterator = iter(dataset) batch = data_iterator.get_next() batch = static_dataset.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) return batch def plot_fixed(garment, body, faces, camera, colors=None): mesh = np.concatenate((body, garment), axis=0) faces = np.concatenate((F_body, faces + 14475), axis=0).astype(np.int32) if colors is None: colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(garment, dtype=np.int32)*230), axis=0).astype(np.float32) / 255.0 else: colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100 / 255.0, colors), axis=0) fig = go.Figure() fig.add_trace(go.Mesh3d( x=mesh[:,0], y=mesh[:,2], z=mesh[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], text="", vertexcolor=colors, lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 }, ),) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),) fig.update_layout(scene = helper) fig.update_layout(scene_camera=camera, width=800, height=1100, autosize=False, margin=dict(l=0, r=0, b=0, t=0, pad=0) ) return fig def custom_imshow(img): # fig = ex.imshow(np.flip(img, axis=0)) fig = ex.imshow(img) fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0), width=img.shape[1], height=img.shape[0]) fig.update_xaxes(visible=False) fig.update_yaxes(visible=False) return fig def color_labels(img): colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, img.max()+1)]) / 255.0 color_img = np.zeros((img.shape[0], img.shape[1], 3), np.float32) for i in range(img.max()+1): color_img[img == i, :] = colors[i] return color_img def color_labels_mesh(img): colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, img.max()+1)]) / 255.0 color_img = np.zeros((img.shape[0], 3), np.float32) for i in range(img.max()+1): color_img[img == i, :] = colors[6-i] return color_img def curvature_sampling(model: CVAECurv, batch, nr=4): 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_mask'], [nr,1,1])) img = np.concatenate(out.numpy(), axis=1) img = color_labels(img.argmax(-1)) custom_imshow(img).show() return out def convert_to_normal_uv(mesh, mask): uv = (convert_to_uv_api(np.squeeze(vertex_normals(mesh, tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4 return uv def get_posed_mesh(uv, pose_shape, body, mask): shape = pose_shape[:, 72:] pose = pose_shape[:, :72] 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])]) if len(uv.shape) == 3: mesh = uv2mesh(uv[tf.newaxis,:,:,:]) else: mesh = uv2mesh(uv) mesh_shaped = mesh + shape_offsets * mask[tf.newaxis,:,tf.newaxis] mesh_posed = smpl.run_garment(pose, shape, mesh_shaped) # mesh_fixed = fix_collision(mesh_posed, body, mask[np.newaxis,:])[0].numpy() # return mesh_fixed return mesh_posed def mm_norm(x): return (x - x.min()) / (x.max() - x.min()) def sample_with_curvature(model: CVAE, batch, sampled_curvature, nr=4): noise = tf.random.normal(shape=(1, 128)) noise = tf.tile(noise, [nr,1]) out_sampled = model.sample(tf.tile(batch['uv_static'], [nr,1,1,1]), tf.tile(batch['uv_body_posed'], [nr,1,1,1]), sampled_curvature, tf.tile(batch['uv_mask'], [nr,1,1]), noise) out_orig = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask'], noise[:1]) # out_sampled, _, _ = 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]), sampled_curvature, tf.tile(batch['uv_mask'], [nr,1,1])], training=False) # out_orig, _, _ = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) out = tf.concat((out_orig, out_sampled), axis=0) mask = batch['mesh_mask'][0].numpy().astype(bool) 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)] plot_fixed(posed_mesh[0][0].numpy()[mask, :], batch['mesh_body_posed'][0].numpy(), faces, camera_1).show() normal_uvs = np.concatenate([convert_to_normal_uv(mesh, mask) for mesh in posed_mesh], axis=1) out_img = np.concatenate(out.numpy(), axis=1) curv_img = color_labels(np.concatenate(np.concatenate((batch['uv_curvature'].numpy(), sampled_curvature), axis=0), axis=1).argmax(-1)) img = np.concatenate((curv_img, normal_uvs), axis=0) custom_imshow(img).show() print() def curv_output_samples(frames): out = [] for frame in frames: batch = get_temp_batch(frame, config_main) mask = batch['mesh_mask'][0].numpy().astype(bool) # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) curv_img = color_labels(batch['uv_curvature'][0].numpy().argmax(-1)) pred, _, _ = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) pred_img = convert_to_normal_uv(uv2mesh(pred), mask) img = np.concatenate((curv_img, pred_img), axis=0) # custom_imshow(img).show() out.append(img) return custom_imshow(np.concatenate(out, axis=1)) def curv_output_samples_mesh(frames): out = [] for frame in tqdm.tqdm(frames): batch = get_temp_batch(frame, config_main) mask = batch['mesh_mask'][0].numpy().astype(bool) _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) curv_colors = color_labels_mesh(batch['mesh_curvature'][0].numpy() + np.ones((14475), np.int32) * mask) pred, _, _ = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) # pred_mesh = uv2mesh(pred) pred_mesh = get_posed_mesh(pred, batch['pose_shape'], batch['mesh_body_posed'], mask) body = batch['mesh_body_posed'][0].numpy() pred_mesh_fixed = postprocess(pred_mesh, body, faces, mask, do_smoothing=False) fig = plot_fixed(pred_mesh_fixed, np.zeros_like(body), faces, camera_1) fig.write_image('/tmp/exp_temp_fasdsavcs.png') img = cv2.imread('/tmp/exp_temp_fasdsavcs.png') fig = plot_fixed(pred_mesh_fixed, np.zeros_like(body), faces, camera_1, curv_colors[mask, :]) fig.write_image('/tmp/exp_temp_fasdsavcs.png') img_curv = cv2.imread('/tmp/exp_temp_fasdsavcs.png') custom_imshow(np.concatenate((img[:,100:-100,:], img_curv[:,100:-100,:]), axis=1)).write_image(os.path.join(out_folder, frame + '_curv.png')) # print() # return custom_imshow(np.concatenate(out, axis=1)) def get_closest_n_points(verts, idx, n): dist = sp.spatial.distance.cdist(verts[idx, :], verts) dist[dist == 0] = 1e9 idx_sort = np.argsort(dist, axis=-1) closest = idx_sort[:, :n] return closest def postprocess(mesh, body, faces, mask, do_smoothing=True): mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :] if do_smoothing: m = tm.Trimesh(mesh_fixed, faces, process=False, validate=False) outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0 outliers = np.nonzero(outliers)[0] outliers_n_neighbours = get_closest_n_points(mesh_fixed, outliers, 2) outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1) faces_fixed = np.concatenate((faces, outliers_faces), axis=0) m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False) tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4) mesh = m.vertices else: mesh = mesh_fixed return mesh def sampling_example(frame, camera, nr=15): batch = get_temp_batch(frame, config_main) mask = batch['mesh_mask'][0].numpy().astype(bool) _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) # 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) # 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=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] # # row_sampling = [] # # fig = plot_fixed(gt_mesh, body, faces, camera) # # fig.write_image('/tmp/exp_temp_7.png') # # img = cv2.imread('/tmp/exp_temp_7.png') # # row_sampling.append(img[:, 100:700, :]) # # for mesh in posed_mesh_fixed: # # fig = plot_fixed(mesh, body, faces, camera) # # fig.write_image('/tmp/exp_temp_7.png') # # img = cv2.imread('/tmp/exp_temp_7.png') # # row_sampling.append(img[:, 100:700, :]) # 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, 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, 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)) # np.save(os.path.join(out_folder, frame, frame + '_gt.npy'), gt_mesh) # np.save(os.path.join(out_folder, frame, frame + '_body.npy'), body) # np.save(os.path.join(out_folder, frame, frame + '_faces.npy'), faces.astype(np.int32)) # for i in range(nr): # np.save(os.path.join(out_folder, frame, frame + '_' + str(i) + '.npy'), posed_mesh_fixed[i]) save_stl(os.path.join(out_folder, frame + '_gt_body.stl'), body, F_body) save_stl(os.path.join(out_folder, frame + '_posed_body.stl'), body, F_body) # custom_imshow(np.concatenate(row_sampling, axis=1)).write_image(os.path.join(out_folder, frame + '_sampling.png')) print() def sampling_poses(frame, camera): nr = 5 frame_nr = [frame.split('_')[1], '25', '35', '50', '0'] bodies = np.stack([np.load(os.path.join(folder, 'uv_posed_body_no_orient_2', frame.split('_')[0] + '_' + f_nr + '.npz'))['uv_posed_body_no_orient'] for f_nr in frame_nr]) bodies_mesh = [np.load(os.path.join(folder, 'mesh_posed_body_no_orient', frame.split('_')[0] + '_' + f_nr + '.npz'))['mesh_posed_body_no_orient'] for f_nr in frame_nr] batch = get_temp_batch(frame, config_main) mask = batch['mesh_mask'][0].numpy().astype(bool) _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) noise = tf.random.normal(shape=(1, 128)) noise = tf.tile(noise, [nr,1]) out = model.sample(tf.tile(batch['uv_static'], [nr,1,1,1]), bodies, tf.tile(batch['uv_curvature'], [nr,1,1,1]), tf.tile(batch['uv_mask'], [nr,1,1]), noise) # 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)] # body = batch['mesh_body_posed'][0].numpy() posed_mesh_fixed = [postprocess(mesh, body, faces, mask, do_smoothing=False) for mesh, body in zip(posed_mesh, bodies_mesh)] row_sampling = [] for mesh, body in zip(posed_mesh_fixed, bodies_mesh): fig = plot_fixed(mesh, body, faces, camera) fig.write_image('/tmp/exp_temp_7.png') img = cv2.imread('/tmp/exp_temp_7.png') row_sampling.append(img[:, 100:700, :]) custom_imshow(np.concatenate(row_sampling, axis=1)).show() print() def sampling_garment(frame, camera): nr = 5 frame_nr = [frame.split('_')[0], '03265', '02869', '00300', '01763'] templates = np.stack([np.load(os.path.join(folder, 'uv_static', f_nr + '.npz'))['uv_static'] for f_nr in frame_nr]) # bodies_mesh = [np.load(os.path.join(folder, 'mesh_posed_body_no_orient', frame.split('_')[0] + '_' + f_nr + '.npz'))['mesh_posed_body_no_orient'] for f_nr in frame_nr] masks = np.stack([np.load(os.path.join(folder, 'mesh_mask', f_nr + '.npy'))['uv_static'] for f_nr in frame_nr]) batch = get_temp_batch(frame, config_main) mask = batch['mesh_mask'][0].numpy().astype(bool) _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) noise = tf.random.normal(shape=(1, 128)) noise = tf.tile(noise, [nr,1]) out = model.sample(templates, 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) # 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)] body = batch['mesh_body_posed'][0].numpy() posed_mesh_fixed = [postprocess(mesh, body, faces, mask, do_smoothing=False) for mesh, mask in zip(posed_mesh, masks)] row_sampling = [] for mesh in posed_mesh_fixed: fig = plot_fixed(mesh, body, faces, camera) fig.write_image('/tmp/exp_temp_7.png') img = cv2.imread('/tmp/exp_temp_7.png') row_sampling.append(img[:, 100:700, :]) custom_imshow(np.concatenate(row_sampling, axis=1)).show() print() if __name__ == '__main__': parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) smpl = SMPL() if config_main['is_skirt']: F = np.load("custom_utils/smpl_faces_skirt.npy") folder = os.path.join('/home/workspace/local-data_128/train_skirt') else: F = np.load("custom_utils/smpl_faces.npy") folder = os.path.join('/home/workspace/local-data/train') # F = np.load("custom_utils/smpl_faces_skirt.npy") F_body = np.load("custom_utils/smpl_faces.npy") # folder = os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if is_skirt else 'test_ready')) # folder = os.path.join('/home/workspace/local-data/train_skirt') model = CVAE(config_main) # curv_model = CVAECurv(config_curv) 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' # faces = np.load(os.path.join(faces_folder, '01763' + '.npy')) frame = '00761_150' body_frames = ['01763_259' , '007292_282'] 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() # 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)) # sample_with_curvature(model, input_batch, sampled_curvature) # out_folder = '/home/workspace/data/results/curv_pred_plots_skirt' # out_folder = '/home/workspace/local-data/results/curv_pred_plots_skirt' # os.makedirs(out_folder, exist_ok=True) # curv_output_samples_mesh(frames) # frame = '01082_109' # 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', '02276_232', '02978_242', '04008_295', '04173_193'] # frames = ['00073_160', '00073_217', '00096_86', \ # '04347_100', '04360_2', '04404_41', '04442_85', '04466_59', '04620_68', '04620_129', \ # '04655_28', '04668_255', '04682_96', '04709_37', '06117_123', '06141_162'] # frames = ['00265_257', '04070_263', '00239_9', '04608_170', '04372_289', '00077_106', '00205_164', '00323_97', '00185_7', '00079_250', \ # '04143_56', '00240_136', '00185_33'] frames = ['009042_205', '009152_33', '010352_168', '012002_4', '03316_168', '03365_230'] out_folder = 'examples/qualitative/bodies3' 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) print()