import os os.environ["CUDA_VISIBLE_DEVICES"] = "5" 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 sys import time import cv2 import random import tqdm import numpy as np import trimesh as tm import cmapy as cm import scipy as sp import plotly.graph_objects as go import plotly.express as ex 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_2_pyr_32_64_128_256_512 import config as config_regular from configs.config_abl_2_pyr_32_64_128_256_512_skirt import config as config_skirt 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 plot_fixed(garment, faces, camera, body=None): if body is not None: mesh = np.concatenate((body, garment), axis=0) faces = np.concatenate((F, faces + 14475), axis=0).astype(np.int32) 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: mesh = garment colors = np.ones_like(garment, dtype=np.float32)*230.0 / 255.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 get_faces(): faces_folder = '/home/workspace/data/results/support/faces' ids = os.listdir(faces_folder) ids = [id.split('.')[0] for id in ids] faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids } faces_folder = '/home/workspace/data/results/support/faces_skirt' ids = os.listdir(faces_folder) ids = [id.split('.')[0] for id in ids] faces_skirt = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids } return { **faces, **faces_skirt } def get_bodies(frames, skirt_mods): folder_base = '/home/workspace/data/dataset' out = dict() for frame in frames: id = frame.split('_')[0] body = np.load(os.path.join((data_folder_skirt if skirt_mods[id] == '_skirt' else data_folder), 'mesh_posed_body_no_orient', frame + '.npz'))['mesh_posed_body_no_orient'] out[frame] = body return out def postprocess(garment, body, id, is_skirt): mask = np.load(os.path.join((data_folder_skirt if is_skirt else data_folder), 'mesh_mask', id + '.npy')) garment_temp = np.zeros((1,14475,3), np.float32) garment_temp[0,mask,:] = garment garment_fixed = fix_collision(garment_temp, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :] return garment_fixed def deorient(mesh, frame, is_skirt): smpl = SMPL() pose = np.load(os.path.join((data_folder_skirt if is_skirt else data_folder), 'pose_shape', frame + '.npy'))[:72].astype(np.float32) rodr = smpl.run_rodrigues(pose[np.newaxis, :]).numpy()[0].astype(np.float32) r = Rotation.from_matrix(rodr[0, :, :]) out = r.apply(mesh, inverse=True) return out def get_temp_batch(frame, config, is_skirt): dataset_resolutions = config['pyramid']['resolutions'].copy() if 512 not in dataset_resolutions: dataset_resolutions.append(512) folder = (data_folder_skirt if is_skirt else data_folder) with open(os.path.join(folder, 'frames_temp.txt'), 'w') as f: f.write(frame) factory = DatasetFactory() dataset = factory.get_dataset(folder=folder, batch_size=1, signature=config['signature_dynamic'], resolutions=dataset_resolutions, 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'], resolutions=dataset_resolutions, 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'])) levels_input, conds = group_inputs_for_levels(batch, config) return levels_input, conds def convert_to_mesh(uvs, is_skirt, mask): uv2mesh = UV2Mesh() uv2mesh_skirt = UV2Mesh(is_skirt=True) out = [] for i in range(len(uvs)): if uvs[i].shape[1] != 512: uvs[i] = tf.image.resize(uvs[i], (512, 256), tf.image.ResizeMethod.LANCZOS5) if is_skirt: mesh = uv2mesh_skirt(uvs[i]) else: mesh = uv2mesh(uvs[i]) out.append(mesh[0].numpy()[mask, :]) return out def get_final_mesh(mesh, frame, is_skirt, body, mask, faces, do_postprocessing=False): smpl = SMPL() pose_shape = np.load(os.path.join((data_folder_skirt if is_skirt else data_folder), 'pose_shape', frame + '.npy')).astype(np.float32) shape = pose_shape[72:][np.newaxis,:] pose = pose_shape[:72][np.newaxis,:] pose[:,:3] = 0 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])]) garment_temp = np.zeros((1,14475,3), np.float32) garment_temp[0,mask,:] = mesh mesh_shaped = garment_temp + shape_offsets * mask[tf.newaxis,:,tf.newaxis] mesh_posed = smpl.run_garment(pose, shape, mesh_shaped) if do_postprocessing: mesh_fixed = fix_collision(mesh_posed, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :] 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) mesh_posed = m.vertices else: mesh_posed = mesh_posed[0].numpy()[mask, :] return mesh_posed 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 get_experiment_plots(frames, skirt_mods, bodies, faces, masks, cameras): model: Pyramid = Pyramid(config_regular) load_model(model, 'checkpoints/' + config_regular['name']) model_skirt: Pyramid = Pyramid(config_skirt) load_model(model_skirt, 'checkpoints/' + config_skirt['name']) result = [] for frame in frames: id = frame.split('_')[0] is_skirt = skirt_mods[id] == '_skirt' config = config_regular if not is_skirt else config_skirt levels_input, conds = get_temp_batch(frame, config, is_skirt) if is_skirt: out = model_skirt([levels_input, conds], training=False) else: out = model([levels_input, conds], training=False) uvs = [out[i][0] for i in range(len(out))] meshes = convert_to_mesh(uvs, is_skirt, masks[id]) final_mesh = get_final_mesh(meshes[-1], frame, is_skirt, bodies[frame], masks[id], faces[id], do_postprocessing=False) final_mesh_fixed = get_final_mesh(meshes[-1], frame, is_skirt, bodies[frame], masks[id], faces[id], do_postprocessing=True) frame_out = [] for mesh in meshes: fig = plot_fixed(mesh, faces[id], cameras[frame]) fig.write_image('/tmp/exp_temp.png') img = cv2.imread('/tmp/exp_temp.png') frame_out.append(img[:, 50:-50, :]) # fig = plot_fixed(final_mesh, faces[id], cameras[frame], bodies[frame]) # fig.write_image('/tmp/exp_temp.png') # img = cv2.imread('/tmp/exp_temp.png') # frame_out.append(img[:, 100:700, :]) fig = plot_fixed(final_mesh_fixed, faces[id], cameras[frame], bodies[frame]) fig.write_image('/tmp/exp_temp.png') img = cv2.imread('/tmp/exp_temp.png') frame_out.append(img[:, 50:-50, :]) result.append(np.concatenate(frame_out, axis=1)) return result def get_gt_image(frames, bodies, faces, cameras, skirt_mods): out = [] for frame in frames: id = frame.split('_')[0] mask = np.load(os.path.join((data_folder_skirt if skirt_mods[id] == '_skirt' else data_folder), 'mesh_mask', id + '.npy')) gt = np.load(os.path.join((data_folder_skirt if skirt_mods[id] == '_skirt' else data_folder), 'mesh_posed_no_orient', frame + '.npz'))['mesh_posed_no_orient'].astype(np.float32) fig = plot_fixed(gt[mask,:], faces[id], cameras[frame], bodies[frame]) fig.write_image('/tmp/exp_temp.png') img = cv2.imread('/tmp/exp_temp.png') gt_unposed = np.load(os.path.join((data_folder_skirt if skirt_mods[id] == '_skirt' else data_folder), 'mesh_unposed_no_shape', frame + '.npz'))['mesh_unposed_no_shape'].astype(np.float32) fig = plot_fixed(gt_unposed[mask,:], faces[id], cameras[frame]) fig.write_image('/tmp/exp_temp.png') img_unposed = cv2.imread('/tmp/exp_temp.png') out.append(np.concatenate((img[:, 50:-50, :], img_unposed[:, 50:-50, :]), axis=1)) return out def get_masks(frames): out = dict() for frame in frames: id = frame.split('_')[0] mask = np.load(os.path.join((data_folder_skirt if skirt_mods[id] == '_skirt' else data_folder), 'mesh_mask', id + '.npy')) out[id] = mask return out if __name__ == '__main__': # uv2mesh = UV2Mesh() # smpl = SMPL() F = np.load("custom_utils/smpl_faces.npy") F_skirt = np.load("custom_utils/smpl_faces_skirt.npy") # garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item() garment_type_table = np.load('custom_data/garment_type_table_train.npy', allow_pickle=True).item() frames = ['05656_221', '00300_69', '06955_117'] # frames = ['02184_157', '02241_142', '06955_117', '06804_192'] # frames = ['01763_259', '007292_282', '00539_296', '00075_4'] # frames = [ '01763_259'] 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)) camera_2 = dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=-0.05), eye=dict(x=-0.9, y=0.7, z=0.05)) # cameras = {'01763_259': camera_1, '007292_282': camera_1, '00539_296': camera_1, '00075_4': camera_2} cameras = {'05656_221': camera_1, '00300_69': camera_1, '06955_117': camera_1} # cameras = {'02184_157': camera_1, '02241_142': camera_1, '06955_117': camera_1, '06804_192': camera_1} # abl_all replace abl_norm maybe experiment = 'abl_2_pyr_32_64_128_256_512' skirt_mods = { frame.split('_')[0]: ('_skirt' if (garment_type_table[frame.split('_')[0]] == 'Skirt' or garment_type_table[frame.split('_')[0]] == 'Dress') else '') for frame in frames } data_folder = '/home/workspace/local-data/train' data_folder_skirt = '/home/workspace/local-data/train_skirt' bodies = get_bodies(frames, skirt_mods) masks = get_masks(frames) faces = { frame.split('_')[0]: _maskFaces(np.zeros((14475,3)), F_skirt if skirt_mods[frame.split('_')[0]] == '_skirt' else F, masks[frame.split('_')[0]])[1] for frame in frames } gts = get_gt_image(frames, bodies, faces, cameras, skirt_mods) results_folder_base = '/home/workspace/data/results' output_folder = os.path.join(results_folder_base, 'ablation_plots') os.makedirs(output_folder, exist_ok=True) images = get_experiment_plots(frames, skirt_mods, bodies, faces, masks, cameras) rows = [np.concatenate((gt, exp), axis=1)[75:-100] for gt, exp in zip(gts, images)] final = np.concatenate(rows, axis=0) custom_imshow(final).show() print()