304 lines
11 KiB
Python
304 lines
11 KiB
Python
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() |