This commit is contained in:
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
+15 -17
View File
@@ -24,7 +24,7 @@ from pathlib import Path
from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl
from utils.utils import save_stl, create_temp_frames
from utils.checkpoints import CheckpointSaver, load_model
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
@@ -38,7 +38,6 @@ from torch.utils.data import DataLoader
import trimesh as tm
from zipfile import ZipFile
from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2
@@ -76,15 +75,13 @@ def step_pyramid(model, batch, config):
return out
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
if __name__ == '__main__':
results_folder = Path('/data-local/127_2/results')
out_folder = Path('temp/pyramid_ablation')
iteration = {'nonskirt': 190000, 'skirt': 200000 }
iteration = {'nonskirt': 250000, 'skirt': 200000 }
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
dataset_resolutions = []
@@ -99,8 +96,9 @@ if __name__ == '__main__':
model['nonskirt'].eval()
model['skirt'].eval()
frames = ['00054_1_Dress', '01070_13_Jumpsuit', '03849_27_Dress', '06057_152_Dress',
frames = ['01070_13_Jumpsuit', '00054_1_Dress', '03849_27_Dress', '06057_152_Dress',
'03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
frames = ['00027_1_Dress']
frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
frames = [Frame(frame_str=frame) for frame in frames]
@@ -109,8 +107,8 @@ if __name__ == '__main__':
config = { 'nonskirt': config_1, 'skirt': config_2 }
for frame in tqdm(frames):
create_temp_frames(config_1['train_dir'], frame)
dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
create_temp_frames(config_1['test_dir'], frame)
dataset = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'
@@ -136,15 +134,15 @@ if __name__ == '__main__':
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
# save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
# fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
# ])
# fig.update_layout(scene=dict(aspectmode='cube'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
])
fig.update_layout(scene=dict(aspectmode='cube'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.show()
pass
pass