371 lines
15 KiB
Python
371 lines
15 KiB
Python
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 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, 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
|
|
from models.pyramid_model_optimized import PyramidNew as Pyramid
|
|
|
|
from configs.iccv.config_abl_pyr_32_64_128_256_512 import config as config_regular
|
|
from configs.iccv.config_abl_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.BILINEAR)
|
|
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/iccv/' + config_regular['name'])
|
|
model_skirt: Pyramid = Pyramid(config_skirt)
|
|
load_model(model_skirt, 'checkpoints/iccv/' + 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, :])
|
|
|
|
for i, mesh in enumerate(meshes):
|
|
save_stl(os.path.join(output_folder, f'{frame}_{i+1}.stl'), mesh, faces[id].astype(np.int32))
|
|
save_stl(os.path.join(output_folder, f'{frame}_final.stl'), final_mesh_fixed, faces[id].astype(np.int32))
|
|
|
|
# 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')
|
|
|
|
save_stl(os.path.join(output_folder, f'{frame}_gt.stl'), gt_unposed[mask,:], faces[id].astype(np.int32))
|
|
save_stl(os.path.join(output_folder, f'{frame}_gt_unposed.stl'), gt[mask,:], faces[id].astype(np.int32))
|
|
save_stl(os.path.join(output_folder, f'{frame}_body.stl'), bodies[frame], F)
|
|
|
|
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()
|
|
|
|
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))
|
|
|
|
frames = ['01070_13']
|
|
cameras = {'01070_13': camera_1}
|
|
# 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']
|
|
# 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_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_128/train_skirt'
|
|
|
|
results_folder_base = 'examples/pyramid'
|
|
|
|
output_folder = os.path.join(results_folder_base)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
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)
|
|
|
|
|
|
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() |