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

save

Browse Source
This commit is contained in:
luni
2023-03-07 08:59:20 +00:00
parent 0b518e7f0e
commit e192f6f9e3
21 changed files with 1004 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plot_ablation_pyramid.py
+29 -18
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "5"
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
import tensorflow as tf
@@ -30,15 +30,15 @@ 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.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 Pyramid
from models.pyramid_model_optimized import PyramidNew as 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
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):
@@ -182,7 +182,7 @@ def convert_to_mesh(uvs, is_skirt, mask):
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)
uvs[i] = tf.image.resize(uvs[i], (512, 256), tf.image.ResizeMethod.BILINEAR)
if is_skirt:
mesh = uv2mesh_skirt(uvs[i])
else:
@@ -239,9 +239,9 @@ def get_closest_n_points(verts, idx, n):
def get_experiment_plots(frames, skirt_mods, bodies, faces, masks, cameras):
model: Pyramid = Pyramid(config_regular)
load_model(model, 'checkpoints/' + config_regular['name'])
load_model(model, 'checkpoints/iccv/' + config_regular['name'])
model_skirt: Pyramid = Pyramid(config_skirt)
load_model(model_skirt, 'checkpoints/' + config_skirt['name'])
load_model(model_skirt, 'checkpoints/iccv/' + config_skirt['name'])
result = []
for frame in frames:
id = frame.split('_')[0]
@@ -267,6 +267,10 @@ def get_experiment_plots(frames, skirt_mods, bodies, faces, masks, cameras):
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')
@@ -297,6 +301,10 @@ def get_gt_image(frames, bodies, faces, cameras, skirt_mods):
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
@@ -322,33 +330,36 @@ if __name__ == '__main__':
# 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 = ['05656_221', '00300_69', '06955_117']
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']
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 = {'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'
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/train_skirt'
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)
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)