import os import sys os.environ["CUDA_VISIBLE_DEVICES"] = '6' 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 tqdm import shutil import random import numpy as np import trimesh as tm import plotly.express as ex import plotly.graph_objects as go from custom_utils.custom_utils import _maskFaces, save_stl from custom_utils.smpl import SMPL from plotly.subplots import make_subplots from scipy.spatial.transform import Rotation as R def plot_fixed(mesh, faces, colors): 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) camera = 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) ) 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 plot_fixed_all(gt, cape, ours, error_cape, error_ours, faces, colors): fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ], horizontal_spacing=0) fig.add_trace(go.Mesh3d( x=gt[:,0], y=gt[:,2], z=gt[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], text='GT', vertexcolor=colors, lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 }, ), row=1, col=1) fig.add_trace(go.Mesh3d( x=cape[:,0], y=cape[:,2], z=cape[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], text=str(error_cape), vertexcolor=colors, lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 }, ), row=1, col=2) fig.add_trace(go.Mesh3d( x=ours[:,0], y=ours[:,2], z=ours[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], text=str(error_ours), vertexcolor=colors, lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 }, ), row=1, col=3) fig.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)) 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, scene2 = helper, scene3 = helper) camera = dict( up=dict(x=0, y=0, z=1), center=dict(x=0, y=-0.3, z=-0.05), eye=dict(x=0.5, y=1.2, z=0.05) ) fig.update_layout(scene_camera=camera, width=2000, height=1100, autosize=False, margin=dict(l=0, r=0, b=0, t=0, pad=0), scene2_camera=camera, scene3_camera=camera) return fig def custom_imshow(img): fig = ex.imshow(np.flip(img, axis=0)) 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 read_obj(path): with open(path, 'r') as f: obj = tm.exchange.obj.load_obj(f, skip_materials=True, maintain_order=True) verts = obj['vertices'] return verts def read_frames_file(file): with open(file, 'r') as f: lines = f.readlines() frames = ['_'.join(line.strip().split('.')[:-1]) for line in lines] # frames_lookup = dict() return frames def get_cape_errors(file, frames): out = dict() errors = np.load(file).mean(-1) for frame, error in zip(frames, errors): out[frame] = error return out def get_errors(file): out = dict() with open(file, 'r') as f: for line in f: frame, error = line.split(' ') out[frame] = float(error) return out def read_gt(folder, frame): split = frame.split('_') file = '_'.join(split[:-1]) + '.' + split[-1] + '.npz' data = np.load(os.path.join(folder, file))['mesh_posed'] return data def add_global_orientation(mesh, frame, folder, inverse=False): split = frame.split('_') file = '_'.join(split[:-1]) + '.' + split[-1] + '.npy' pose_shape = np.load(os.path.join(folder, file)) rodr = smpl.run_rodrigues(pose_shape[:72][np.newaxis,:].astype(np.float32)).numpy() r = R.from_matrix(rodr[0, 0, :, :]) out = r.apply(mesh, inverse=inverse) return out if __name__ == '__main__': folder_f = 'CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_female/test_reconstruction_objs_dataset_female_4clotypes' folder_m = 'CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_male/test_reconstruction_objs_dataset_male_4clotypes' frames_file_f = '/home/workspace/data/dataset/cape_processed/datasets/dataset_female_4clotypes_frames.txt' frames_file_m = '/home/workspace/data/dataset/cape_processed/datasets/dataset_male_4clotypes_frames.txt' frames_f = read_frames_file(frames_file_f) frames_f_to_i = { frame: i for i, frame in enumerate(frames_f) } errors_cape_f = get_cape_errors('CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_female/errors.npy', frames_f) errors_cape_f = dict(reversed(sorted(errors_cape_f.items(), key=lambda item: item[1]))) frames_m = read_frames_file(frames_file_m) frames_m_to_i = { frame: i for i, frame in enumerate(frames_m) } errors_cape_m = get_cape_errors('CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_male/errors.npy', frames_m) errors_cape_m = dict(reversed(sorted(errors_cape_m.items(), key=lambda item: item[1]))) errors_f = get_errors('temp/archive/error_abl_2_pyr_64_128_256_512_cape_f.txt') errors_m = get_errors('temp/archive/error_abl_2_pyr_64_128_256_512_cape_m.txt') folder_pyr_f = '/home/workspace/data/results/cape/abl_2_pyr_64_128_256_512_cape_f' folder_pyr_m = '/home/workspace/data/results/cape/abl_2_pyr_64_128_256_512_cape_m' gt_folder_f = '/home/workspace/local-data/cape_processed/test/female/mesh_posed' gt_folder_m = '/home/workspace/local-data/cape_processed/test/male/mesh_posed' pose_folder_m = '/home/workspace/local-data/cape_processed/test/male/pose_shape' pose_folder_f = '/home/workspace/local-data/cape_processed/test/female/pose_shape' smpl = SMPL(orig_smpl=True) # verts_idx = smpl.to_keep clothing_verts_idx = np.load(os.path.join('CAPE', 'data', 'clothing_verts_idx.npy')) cape_mask = np.zeros((6890), bool) cape_mask[clothing_verts_idx] = True cape_mask = cape_mask[smpl.to_keep] J0 = np.load('custom_data/smpl_root_joint.npy') _, faces = _maskFaces(np.zeros((3638, 3)), np.load('custom_utils/smpl_faces_orig.npy'), cape_mask) faces = faces.astype(np.int32) colors = np.ones((3627, 3), np.int8) * 230 gender = 'female' if gender == 'male': folder_cape = folder_m errors_cape = errors_cape_m folder = folder_pyr_m errors = errors_m frames_to_i = frames_m_to_i gt_folder = gt_folder_m pose_folder = pose_folder_m else: folder_cape = folder_f errors_cape = errors_cape_f folder = folder_pyr_f errors = errors_f frames_to_i = frames_f_to_i gt_folder = gt_folder_f pose_folder = pose_folder_f all_errors_1 = [] all_errors_2 = [] for key in tqdm.tqdm(errors_cape.keys()): # count = 0 # while True: if os.path.exists(f'temp/cape_stl/{key}_gt.stl'): continue # key = random.choice(list(errors_cape.keys())) cape = read_obj(os.path.join(folder_cape, '{:0>4d}.obj'.format(frames_to_i[key])))[clothing_verts_idx, :] - J0[np.newaxis, :] split = key.split('_') file = '_'.join(split[:-1]) + '.' + split[-1] + '.npy' pose_shape_path = os.path.join(pose_folder, file) if not os.path.exists(pose_shape_path): continue cape = add_global_orientation(cape, key, pose_folder, inverse=True) if not os.path.exists(os.path.join(folder, key + '_posed.npy')): continue ours = np.load(os.path.join(folder, key + '_posed.npy'))[cape_mask, :] # ours = add_global_orientation(ours, key, pose_folder) gt = read_gt(gt_folder, key)[cape_mask, :] gt = add_global_orientation(gt, key, pose_folder, inverse=True) # error_1 = np.sqrt(((gt - cape) ** 2).sum(-1)).mean() # error_2 = np.sqrt(((gt - ours) ** 2).sum(-1)).mean() # all_errors_1.append(error_1) # all_errors_2.append(error_2) # print(np.mean(all_errors_1)) # print(np.mean(all_errors_2)) # fig = plot_fixed_all(gt, cape, ours, errors_cape[key], errors[key], faces, colors) # fig.write_image('temp/cape_2/' + key + '.png') save_stl(f'temp/cape_stl/{key}_gt.stl', gt, faces) save_stl(f'temp/cape_stl/{key}_cape.stl', cape, faces) save_stl(f'temp/cape_stl/{key}_ours.stl', ours, faces) # print(count) # count += 1 print()