luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

commit all

Browse Source
This commit is contained in:
luni
2023-01-25 12:25:59 +00:00
parent 1bf224841f
commit 7470ca9fb6
139 changed files with 82069 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plot_ablation_pyramid.py
+360
View File
@@ -0,0 +1,360 @@
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()