243 lines
10 KiB
Python
243 lines
10 KiB
Python
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
|
|
import sys
|
|
import time
|
|
import cv2
|
|
import random
|
|
import tqdm
|
|
import numpy as np
|
|
import trimesh as tm
|
|
import tensorflow as tf
|
|
import cmapy as cm
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
from plotly.subplots import make_subplots
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
import tps
|
|
|
|
from dataset import DatasetFactory, StaticDataset
|
|
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces, load_model, save_stl
|
|
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
|
|
|
|
|
|
|
|
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, body, faces, camera):
|
|
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
|
|
|
|
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/local-data'
|
|
out = dict()
|
|
for frame in tqdm.tqdm(frames):
|
|
id = frame.split('_')[0]
|
|
body = np.load(os.path.join(folder_base, 'test' + skirt_mods[id] + '_ready' + ('_2' if skirt_mods[id]!='' else '' ), '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('/home/workspace/data/dataset', ('test_skirt_ready_2' if is_skirt else 'test_ready'), '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('/home/workspace/data/dataset', ('test_skirt_ready_2' if is_skirt else 'test_ready'), '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_experiment_plots(folder, frames, experiments, skirt_mods, bodies, faces, cameras):
|
|
out = []
|
|
for frame in tqdm.tqdm(frames):
|
|
frame_out = []
|
|
id = frame.split('_')[0]
|
|
for exp in experiments:
|
|
garment = np.load(os.path.join(folder, exp[0] + skirt_mods[id] + '-' + exp[1], frame + '.npy'))
|
|
if exp == 'abl_orient':
|
|
garment = deorient(garment, frame, skirt_mods[id] == '_skirt')
|
|
garment = postprocess(garment, bodies[frame], id, skirt_mods[id] == '_skirt')
|
|
fig = plot_fixed(garment, bodies[frame], faces[id], cameras[frame])
|
|
fig.write_image('/tmp/exp_temp_5.png')
|
|
img = cv2.imread('/tmp/exp_temp_5.png')
|
|
frame_out.append(img[:, 100:700, :])
|
|
|
|
out.append(np.concatenate(frame_out, axis=1))
|
|
|
|
return out
|
|
|
|
def get_experiment_meshes(folder, frames, experiments, skirt_mods, bodies, faces, cameras):
|
|
out = []
|
|
for frame in tqdm.tqdm(frames):
|
|
frame_out = []
|
|
id = frame.split('_')[0]
|
|
for exp in experiments:
|
|
garment = np.load(os.path.join(folder, exp[0] + skirt_mods[id] + '-' + exp[1], frame + '.npy'))
|
|
if exp == 'abl_orient':
|
|
garment = deorient(garment, frame, skirt_mods[id] == '_skirt')
|
|
garment = postprocess(garment, bodies[frame], id, skirt_mods[id] == '_skirt')
|
|
frame_out.append([garment, bodies[frame], faces[id]])
|
|
|
|
out.append(frame_out)
|
|
return out
|
|
|
|
def get_gt_image(frames, bodies, faces, cameras, skirt_mods):
|
|
out = []
|
|
for frame in tqdm.tqdm(frames):
|
|
id = frame.split('_')[0]
|
|
mask = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_mask', id + '.npy'))
|
|
gt = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_posed_no_orient', frame + '.npz'))['mesh_posed_no_orient'].astype(np.float32)
|
|
fig = plot_fixed(gt[mask,:], bodies[frame], faces[id], cameras[frame])
|
|
fig.write_image('/tmp/exp_temp_5.png')
|
|
img = cv2.imread('/tmp/exp_temp_5.png')
|
|
out.append(img[:, 100:700, :])
|
|
|
|
return out
|
|
|
|
|
|
def get_gt_meshes(frames, bodies, faces, cameras, skirt_mods):
|
|
out = []
|
|
for frame in tqdm.tqdm(frames):
|
|
id = frame.split('_')[0]
|
|
mask = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_mask', id + '.npy'))
|
|
gt = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_posed_no_orient', frame + '.npz'))['mesh_posed_no_orient'].astype(np.float32)
|
|
out.append([gt[mask,:], bodies[frame], faces[id]])
|
|
|
|
return out
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# uv2mesh = UV2Mesh()
|
|
# smpl = SMPL()
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
|
|
|
|
faces = get_faces()
|
|
|
|
frames = ['00079_135']
|
|
# frames = [ '00075_4', '00475_273', '00761_295', '00793_3', '00993_220', '01134_112', '01278_108', '00718_38', '00951_34', '00962_60','01065_38','01191_83','01240_168']
|
|
# rand_frames = [frame.split('.')[0] for frame in os.listdir('/home/workspace/data/dataset/test_ready/pose_shape')]
|
|
# random.shuffle(rand_frames)
|
|
# rand_frames = rand_frames[:25]
|
|
# frames.extend(rand_frames)
|
|
# rand_frames = [frame.split('.')[0] for frame in os.listdir('/home/workspace/data/dataset/test_skirt_ready_2/pose_shape')]
|
|
# random.shuffle(rand_frames)
|
|
# rand_frames = rand_frames[:25]
|
|
# frames.extend(rand_frames)
|
|
|
|
|
|
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 = {frame: camera_1 for frame in frames}
|
|
|
|
# abl_all replace abl_norm maybe
|
|
experiments = [['iccv_pyr_32', '335000'], ['iccv_pyr_32_skin', '330000'], ['iccv_pyr_32_curv', '155000'], ['iccv_pyr_32_reg', '150000']] # , 'pyra_32_res'
|
|
skirt_mods = { frame.split('_')[0]: ('_skirt' if os.path.exists('/home/workspace/data/dataset/test_skirt_ready_2/pose_shape/' + frame + '.npy') else '') for frame in frames }
|
|
|
|
bodies = get_bodies(frames, skirt_mods)
|
|
# gts = get_gt_image(frames, bodies, faces, cameras, skirt_mods)
|
|
gts = get_gt_meshes(frames, bodies, faces, cameras, skirt_mods)
|
|
|
|
results_folder_base = '/home/workspace/data/results/iccv/'
|
|
|
|
output_folder = os.path.join(results_folder_base, 'ablation_plots')
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
# images = get_experiment_plots(results_folder_base, frames, experiments, skirt_mods, bodies, faces, cameras)
|
|
meshes = get_experiment_meshes(results_folder_base, frames, experiments, skirt_mods, bodies, faces, cameras)
|
|
for mesh, frame in zip(meshes, frames):
|
|
for exp, (garment, body, faces) in zip(experiments, mesh):
|
|
save_stl(f'examples/{frame}_{exp[0]}_garment.stl', garment, faces)
|
|
save_stl(f'examples/{frame}_{exp[0]}_body.stl', body, F)
|
|
|
|
for mesh, frame in zip(gts, frames):
|
|
garment, body, faces = mesh
|
|
save_stl(f'examples/{frame}_gt_garment.stl', garment, faces)
|
|
save_stl(f'examples/{frame}_gt_body.stl', body, F)
|
|
|
|
# rows = [np.concatenate((gt, exp), axis=1) for gt, exp in zip(gts, images)]
|
|
|
|
# for frame, row in zip(frames, rows):
|
|
# custom_imshow(row).write_image(os.path.join(output_folder, frame + '_ablation.png'))
|
|
|
|
# final = np.concatenate(rows, axis=0)
|
|
|
|
# custom_imshow(final).show()
|
|
|
|
print() |