From 37ad3a3f438d9c5cf850752802b492846c673a2b Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 4 Jul 2023 23:52:13 +0200 Subject: [PATCH] remove long edges in post --- custom_utils/custom_utils.py | 35 +++- data/dataset.py | 12 +- data/preprocessing.py | 86 ++++++++- scripts/cvae_pyramid.py | 337 +++-------------------------------- 4 files changed, 151 insertions(+), 319 deletions(-) diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 58bb990..b41c681 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -287,4 +287,37 @@ def write_obj(file, vertices, faces, colors): file.write(line) for f in faces: line = 'f ' + ' '.join([str(_) for _ in f]) + '\n' - file.write(line) \ No newline at end of file + file.write(line) + + +def group_inputs_for_levels(batch, config): + levels_input = [] + for res in config['pyramid']['resolutions']: + res_modifier = '_' + str(res) if res != 512 else '' + if config['use_skinning']: + uv_orig = batch['uv_unposed' + res_modifier] + else: + uv_orig = batch['uv_posed' + res_modifier] + levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_normals' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) + conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_normals']] + return levels_input, conds + + +def remove_long_edges(vertices, faces, limit=0.1): + # Compute edge lengths + edges = np.vstack((faces[:,[0,1]], faces[:,[1,2]], faces[:,[2,0]])) + edge_lengths = np.linalg.norm(vertices[edges[:,0]] - vertices[edges[:,1]], axis=1) + + # Find edges longer than 0.1 + long_edges = edges[edge_lengths > limit] + + # Find vertices that are part of long edges + vertices_to_remove = np.unique(long_edges) + + # Remove vertices and update faces + mask = np.ones(vertices.shape[0], dtype=bool) + mask[vertices_to_remove] = False + new_vertices = vertices[mask] + + + return new_vertices, mask \ No newline at end of file diff --git a/data/dataset.py b/data/dataset.py index 4da75f7..1b38abc 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -80,6 +80,12 @@ class StaticCache(): mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool) data['mesh_mask'] = mesh_mask + mesh_fix = np.load(os.path.join(self.folder, 'mesh_fixed', f'{id}_{garment}.npz')) + data['mesh_fix'] = dict() + data['mesh_fix']['mesh_mask'] = torch.from_numpy(mesh_fix['mask']) + # data['mesh_fix']['mask_convert'] = torch.from_numpy(mesh_fix['mas_convert']) + # data['mesh_fix']['faces'] = torch.from_numpy(mesh_fix['faces']) + for res in self.resolutions: res_string = '' if res != 512: @@ -107,13 +113,15 @@ class StaticCache(): def __getitem__(self, frame: Frame) -> dict: if (frame.id, frame.garment) not in self.store: item = self.read(frame.id, frame.garment) - stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} + stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items() if type(value) == np.ndarray} + stored_data['mesh_fix'] = item['mesh_fix'] for key, value in stored_data.items(): if key.startswith('uv_') and value.shape[-1] == 3: stored_data[key] = torch.permute(value, [2,0,1]) self.store[(frame.id, frame.garment)] = stored_data - data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()} + data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items() if type(value) == torch.Tensor } + data['mesh_fix'] = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)]['mesh_fix'].items() } return data diff --git a/data/preprocessing.py b/data/preprocessing.py index ea09635..5089977 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -33,7 +33,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris, readPC2Frame from data.utils import Frame from smpl.smpl_torch import SMPL -from custom_utils.custom_utils import _maskFaces +from custom_utils.custom_utils import _maskFaces, remove_long_edges def create_extras_file(folder): @@ -631,7 +631,6 @@ def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') - def check_irregular_poses(folder, folder_reg, out_folder): for i in range(100, 200, 5): frame: Frame = Frame('01512', i, 'Dress') @@ -800,6 +799,83 @@ def get_ids(): pass +def generate_fixed_masks(folder, out_folder): + ids = sorted(os.listdir(folder)) + os.makedirs(os.path.join(out_folder, 'mesh_fixed'), exist_ok=True) + F = np.load('smpl/smpl_faces.npy') + F_skirt = np.load('smpl/smpl_faces_skirt.npy') + + import random + # import trimesh + import pyvista as pv + random.shuffle(ids) + + + for id in tqdm(ids): + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + + for garment in garments: + if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): + continue + + is_skirt = garment in ['Dress', 'Skirt'] + + mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy')) + mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) + _, faces = _maskFaces(np.zeros((14475, 3)), F_skirt if is_skirt else F, mask) + _, new_mask = remove_long_edges(mesh_static, faces.astype(np.int32), 0.2 if is_skirt else 0.1) + new_mask_full = np.zeros((14475), dtype=bool) + new_mask_full[mask] = new_mask + _, new_faces = _maskFaces(np.zeros((14475, 3)), F_skirt if is_skirt else F, new_mask_full) + + mesh_new = mesh_static[new_mask] + + # fig = go.Figure(data=[go.Mesh3d(x=mesh_new[:,0], y=mesh_new[:,1], z=mesh_new[:,2], i=new_faces[:,0], j=new_faces[:,1], k=new_faces[:,2])]) + # fig.update_layout(scene_aspectmode='cube') + # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) + # fig.show() + + flattened_faces = np.asarray([np.concatenate((np.array([3]), f)) for f in new_faces.astype(np.int32)]).flatten() + m = pv.PolyData(mesh_new, flattened_faces) + + m.fill_holes(0.03, inplace=True, progress_bar=False) + + # m = trimesh.Trimesh(vertices=mesh_new, faces=new_faces, process=False, validate=False) + # trimesh.repair.fill_holes(m) + + mesh_new = np.asarray(m.points) + new_faces = np.asarray(m.faces).reshape([-1, 4])[:, 1:] + + m = tm.Trimesh(vertices=mesh_new, faces=new_faces, process=False, validate=False) + tm.repair.fix_winding(m) + + mesh_new = np.asarray(m.vertices) + new_faces = np.asarray(m.faces) + + # fig = go.Figure(data=[go.Mesh3d(x=mesh_new[:,0], y=mesh_new[:,1], z=mesh_new[:,2], i=new_faces[:,0], j=new_faces[:,1], k=new_faces[:,2])]) + # fig.update_layout(scene_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=mesh_static[:,0], y=mesh_static[:,1], z=mesh_static[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])]) + # fig.update_layout(scene_aspectmode='cube') + # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) + # fig.show() + + # print((~new_mask).sum()) + + np.savez_compressed(os.path.join(out_folder, 'mesh_fixed', f'{id}_{garment}.npz'), mask=new_mask_full, mas_convert=new_mask, faces=new_faces) + + pass + + + + + if __name__ == '__main__': @@ -832,9 +908,11 @@ if __name__ == '__main__': # calculate_average(folder, out_folder) # create_frames_file(folder, out_folder) - get_ids() + # get_ids() # create_frames_file_test(folder, out_folder) # create_frames_file_partial(folder, out_folder, part=0.5) # folder = '/data-local/127_2/cape_processed/test/female' - # create_mesh_normals_cape(folder, None) \ No newline at end of file + # create_mesh_normals_cape(folder, None) + + generate_fixed_masks(folder, out_folder) \ No newline at end of file diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index 9779384..762b81d 100644 --- a/scripts/cvae_pyramid.py +++ b/scripts/cvae_pyramid.py @@ -238,8 +238,6 @@ def train_step(batch): res_modifier = '_' + str(res) if res != 512 else '' pred, mean, logvar = out_levels[i] - - l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) kl = kl_loss(mean, logvar) @@ -251,16 +249,16 @@ def train_step(batch): gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) - mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None] mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] else: mesh_pred_posed = mesh_pred - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask']) c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) - n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss) if not config['is_cape']: @@ -290,19 +288,19 @@ def train_step(batch): gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) - mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None] mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] else: mesh_pred_posed = mesh_pred - - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask']) c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) - n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) - metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask']) output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl, @@ -343,18 +341,18 @@ def valid_step(batch): gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) - mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_fix']['mesh_mask'][:,:,None] mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] else: mesh_pred_posed = mesh_pred - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_fix']['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_fix']['mesh_mask']) c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) - n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) - + kl = kl_loss(mean, logvar) if config['use_extra_losses']: @@ -373,7 +371,7 @@ def valid_step(batch): logvar = torch.mean(logvar) ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) - metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask']) output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl, @@ -385,63 +383,6 @@ def valid_step(batch): return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output } -def test_step(batch): - levels_input, conds = group_inputs_for_levels(batch, config) - out_levels = model([levels_input, conds], training=False) - - shape = batch['pose_shape'][:, 72:] - gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) - shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, dim=0) for i in range(shape.shape[0])]) - mesh_orig = batch['mesh_unposed'] - - output = dict() - reconstr_loss_total = 0 - kl_total = 0 - - #for i, res in enumerate(config['pyramid']['resolutions']): - i = -1 - res = config['pyramid']['resolutions'][-1] - res_modifier = '_' + str(res) if res != 512 else '' - pred, mean, logvar = out_levels[i] - - mesh_pred = uv2mesh(pred) - mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] - mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) - - # is the output already masked? - - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) - c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) - l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), dim=-1), dim=[1, 2])) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1), dim=-1)) - # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) - # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) - kl = kl_loss(mean, logvar) - - reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss) - if not config['is_cape']: - reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] - # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']) - kl_total += kl - - mean = torch.mean(mean) - logvar = torch.mean(logvar) - - metric = torch.sum(tf.math.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask']) - euclid_error_per_garment = torch.sum(tf.math.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, dim=-1)), dim=1) / torch.sum(batch['mesh_mask'], dim=1) - - output = { **output, **{ 'test/edge_loss' + res_modifier: e_loss, 'test/collision_loss' + res_modifier: c_loss,'test/euclid_error' + res_modifier: metric, - 'test/mean' + res_modifier: mean, 'test/logvar' + res_modifier: logvar, 'test/kl_loss' + res_modifier: kl, - 'test/var' + res_modifier: np.e ** logvar, 'test/l2_loss' + res_modifier: l2_loss, - 'test/euclid_error': metric, 'test/collision_ratio' + res_modifier: c_ratio } } - - loss = reconstr_loss_total + balancing_coeff * kl_total - - return { **{ 'test/loss': loss, 'test/reconstr_loss_total': reconstr_loss_total, 'test/kl_total': kl_total }, **output }, euclid_error_per_garment - - - def plot_everything(batch: dict, type: str): if config['use_skinning']: mesh_orig, uv_orig = batch['mesh_unposed'].cpu().numpy(), batch['uv_unposed'] @@ -473,8 +414,11 @@ def plot_everything(batch: dict, type: str): imgs = [wandb.Image(img[i].cpu().numpy()) for i in range(4)] _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") - mask = batch['mesh_mask'].cpu().numpy().astype(bool) - _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask) + mask = batch['mesh_fix']['mesh_mask'].cpu().numpy().astype(bool) + if 'mesh_fix' in batch: + faces = [np.load( (dataset_train if type == 'train' else dataset_valid).folder + "/mesh_fixed/" + batch['id'][i] + '_' + batch['garment'][i] + ".npz" )['faces'] for i in range(len(batch['id']))] + else: + _, faces = maskFacesBatch(np.ones_like(mesh_orig), _F, mask) if config['use_skinning']: shape = batch['pose_shape'][:, 72:] @@ -543,237 +487,6 @@ def plot_everything(batch: dict, type: str): return { **figs, **figs_coll, **{ type + '/images': imgs, }, **{ f'{type}/batch_info': batch_info }, } -def sample(folder): - # run = wandb.init(project='vae_plots', name=folder.split('/')[-1]) - load_model(model,folder) - count = 0 - for batch in dataset_train: - batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) - nr_of_samples=2 - - pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) - pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] - pose, shape = tf.tile(pose[None,:], [nr_of_samples,1]), tf.tile(shape[None,:], [nr_of_samples,1]) - - mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) - mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy() - mesh_orig = batch['mesh_posed'][0].numpy() - mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0] - - body = batch['mesh_body_posed'][0].numpy() - body_color = np.ones_like(body, dtype=np.int32)*100 - - F = np.load("custom_utils/smpl_faces.npy") - mask = batch['mesh_mask'].numpy().astype(bool)[0] - _, faces = _maskFaces(np.ones((14475, 3)), F, mask) - - - pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True) - mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2])) - mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy() - - - - - meshes = [] - meshes.append(go.Mesh3d( - # 8 vertices of a cube - x=mesh_orig_reposed[mask,0], - y=mesh_orig_reposed[mask,2], - z=mesh_orig_reposed[mask,1], - # i, j and k give the vertices of triangles - i = faces[:,0], - j = faces[:,1], - k = faces[:,2] - )) - - for i in range(nr_of_samples): - x = mesh_posed[i][mask, :].copy() - # x[:, 2] += 1 * (i + 1) - x[:, 2] += 1 - meshes.append(go.Mesh3d( - # 8 vertices of a cube - x=x[:,0], - y=x[:,2], - z=x[:,1], - # i, j and k give the vertices of triangles - i = faces[:,0], - j = faces[:,1], - k = faces[:,2] - )) - - - for i in range(nr_of_samples): - x = mesh_posed_zero[i][mask, :].copy() - # x[:, 2] += 1 * (i + 1) - x[:, 2] += 2 - meshes.append(go.Mesh3d( - # 8 vertices of a cube - x=x[:,0], - y=x[:,2], - z=x[:,1], - # i, j and k give the vertices of triangles - i = faces[:,0], - j = faces[:,1], - k = faces[:,2] - )) - - # for i in range(5): - # x = body.copy() - # x[:, 2] += 1 * i - # meshes.append(go.Mesh3d( - # # 8 vertices of a cube - # x=x[:,0], - # y=x[:,2], - # z=x[:,1], - # # i, j and k give the vertices of triangles - # i = F[:,0], - # j = F[:,1], - # k = F[:,2], - # vertexcolor=body_color - # )) - - - fig = go.Figure(data=meshes) - - fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) - # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1)) - fig.update_layout( - scene = dict( - xaxis = dict(nticks=4, range=[-1,1],), - # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],), - yaxis = dict(nticks=4, range=[-1,3],), - zaxis = dict(nticks=4, range=[-1,1],),) - ) - - # wandb.log({'plot_' + str(count): fig}) - fig.show() - - pred = pred.numpy() - frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])] - frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:], - save_all=True, duration=500, loop=0) - - - if count > 10: - break - - print(count) - count += 1 - - print() - - -def test(): - garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item() - str_split = config['name'].split('_') - test_idx = str_split.index('test') - checkpoint_name = '_'.join(str_split[:test_idx]) - load_model(model, 'checkpoints/' + checkpoint_name) - - dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False) - static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=False) - - step = 0 - losses = [] - - if config['is_skirt']: - per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] } - else: - per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] } - - for batch in tqdm.tqdm(dataset_test): - batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) - - test_loss, euclid_error_per_garment = test_step(batch) - losses.append(test_loss) - - for i, id in enumerate(batch['id']): - id_str = id.numpy().decode("utf-8") - per_garment_errors[garment_type_table[id_str]].append(euclid_error_per_garment[i].numpy()) - - to_log = { **{ 'current_step': step }, **test_loss } - if step % 1000 == 0: - figs = plot_everything(batch, 'test') - to_log = { **to_log, **figs } - - wandb.log(to_log, step = step) - - # print(step) - step += 1 - - # averages = dict() - path_to_plotly_html = "temp/plotly_figure.html" - - histograms = dict() - for key in test_loss.keys(): - accumulated = [losses[i][key].numpy() for i in range(len(losses))] - fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean())) - fig.write_html(path_to_plotly_html, auto_play = False) - - histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html) - # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean() - - for key in per_garment_errors.keys(): - fig = ex.histogram(per_garment_errors[key], title='Average: ' + str(np.asarray(per_garment_errors[key]).mean())) - fig.write_html(path_to_plotly_html, auto_play = False) - histograms['all/' + key + 'euclid_error_hist'] = wandb.Html(path_to_plotly_html) - - - wandb.log({ **histograms }, step = step) - wandb.finish() - pass - - -def generate_results(): - load_model(model, 'checkpoints/' + config['name']) - - dataset_test = factory.get_dataset(folder=config['test_dir'], - batch_size=config['batch_size'], - # batch_size=1, - signature=config['signature_dynamic'], - resolutions=dataset_resolutions, - in_memory=config['dataset_in_memory'], - is_cape=config['is_cape'], - frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, - do_repeat=False, - shuffle=False) - static_dataset_test = StaticDataset(folder=config['test_dir'], - signature=config['signature_static'], - resolutions=dataset_resolutions, - in_memory=True, - is_cape=config['is_cape']) - folder = os.path.join('/home/workspace/data/results/cape', config['name']) - os.makedirs(folder, exist_ok=True) - - for batch in tqdm.tqdm(dataset_test): - batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) - - levels_input, conds = group_inputs_for_levels(batch, config) - out_levels = model([levels_input, conds], training=False) - - shape = batch['pose_shape'][:, 72:] - gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) - shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) - - i = -1 - pred = out_levels[i][0] - - mesh_pred = uv2mesh(pred) - mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] - mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) - - for i in range(batch['id'].shape[0]): - mask = batch['mesh_mask'][i].numpy().astype(bool) - id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8") - unposed = mesh_pred[i].numpy()[mask, :] - posed = mesh_pred_posed[i].numpy()[mask, :] - np.save(os.path.join(folder, id + '_' + frame_nr + '_posed.npy'), posed) - np.save(os.path.join(folder, id + '_' + frame_nr + '_unposed.npy'), unposed) - - pass - def speed_step(levels_input, conds): return model.sample([levels_input, conds], training=False) @@ -840,7 +553,7 @@ if __name__ == '__main__': optimizer = optim.Adam(model.parameters(), lr=config['lr']) _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) - ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=5) + ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20) scaler = torch.cuda.amp.GradScaler(enabled=use_amp) g_balancing_coeff = config['starting_balancing_coeff']