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-04 12:36:00 +00:00
parent 165adf7953
commit 7c92e03b45
28 changed files with 1216 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plot_ablation.py
+55 -19
View File
@@ -17,7 +17,7 @@ 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
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
@@ -92,7 +92,7 @@ def get_faces():
def get_bodies(frames, skirt_mods):
folder_base = '/home/workspace/data/dataset'
folder_base = '/home/workspace/local-data'
out = dict()
for frame in tqdm.tqdm(frames):
id = frame.split('_')[0]
@@ -128,7 +128,7 @@ def get_experiment_plots(folder, frames, experiments, skirt_mods, bodies, faces,
frame_out = []
id = frame.split('_')[0]
for exp in experiments:
garment = np.load(os.path.join(folder, exp + skirt_mods[id] + '-70000', frame + '.npy'))
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')
@@ -141,6 +141,20 @@ def get_experiment_plots(folder, frames, experiments, skirt_mods, bodies, faces,
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 = []
@@ -155,6 +169,17 @@ def get_gt_image(frames, bodies, faces, cameras, skirt_mods):
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()
@@ -165,16 +190,16 @@ if __name__ == '__main__':
faces = get_faces()
frames = []
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)
# 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))
@@ -182,23 +207,34 @@ if __name__ == '__main__':
cameras = {frame: camera_1 for frame in frames}
# abl_all replace abl_norm maybe
experiments = ['abl_norm', 'abl_skin', 'abl_curv', 'abl_extra_loss', 'abl_abs']
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_image(frames, bodies, faces, cameras, skirt_mods)
gts = get_gt_meshes(frames, bodies, faces, cameras, skirt_mods)
results_folder_base = '/home/workspace/data/results'
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)
# 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)
rows = [np.concatenate((gt, exp), axis=1) for gt, exp in zip(gts, images)]
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)
for frame, row in zip(frames, rows):
custom_imshow(row).write_image(os.path.join(output_folder, frame + '_ablation.png'))
# 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)