From 882ddf5578b10e35c56ee444821b504bc0a4480d Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 8 Jun 2023 10:54:34 +0000 Subject: [PATCH] use new colors for pose --- data/pose_dataset.py | 4 ++-- data/pose_renderer.py | 9 ++++++--- scripts/pose_encoder_train.py | 32 +++++++++++++++----------------- 3 files changed, 23 insertions(+), 22 deletions(-) diff --git a/data/pose_dataset.py b/data/pose_dataset.py index 295bf13..99d3951 100644 --- a/data/pose_dataset.py +++ b/data/pose_dataset.py @@ -102,8 +102,8 @@ class Dataset(Dataset): # 'body_out': body_out, # 'render': render, # 'render_out': render_out, - 'mask': mask, - 'mask_out': mask_out, + 'mask': mask, + 'mask_out': mask_out, } if self.cache is not None: diff --git a/data/pose_renderer.py b/data/pose_renderer.py index a78cd8b..4af8b40 100644 --- a/data/pose_renderer.py +++ b/data/pose_renderer.py @@ -30,9 +30,12 @@ from pytorch3d.structures import Meshes class PoseRenderer(Module): def __init__(self, device='cuda') -> None: super().__init__() - colors = np.load('custom_data/pose_colors.npy') - # colors = np.load('custom_data/body_colors_interpolated.npy').astype(np.float32) / 255.0 - + # colors = np.load('custom_data/pose_colors.npy') + colors_ = np.load('custom_data/body_colors_interpolated.npy').astype(np.float32) / 255.0 + color_order = np.load('custom_data/color_ordered_indices.npy') + colors = np.zeros_like(colors_) + colors[color_order] = colors_[np.arange(27578)] + self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :]) self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy() diff --git a/scripts/pose_encoder_train.py b/scripts/pose_encoder_train.py index e55bc03..1023c44 100644 --- a/scripts/pose_encoder_train.py +++ b/scripts/pose_encoder_train.py @@ -94,19 +94,20 @@ def valid_step(batch): symm_rotation = loss_fn_rotation(pred_rotation) loss = mse + mse_mask + symm_rotation - return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation } def plot_everything(batch, type: str): - pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']]) + model.eval() + with torch.inference_mode(): + pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']]) mask = torch.permute(batch['mask'], (0, 2, 3, 1)) mask_out = torch.permute(batch['mask_out'], (0, 2, 3, 1)) pose = torch.permute(batch['pose'], (0, 2, 3, 1)) pose_out = torch.permute(batch['pose_out'], (0, 2, 3, 1)) - pred_mask = torch.permute(pred_mask, (0, 2, 3, 1)).detach() - pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)).detach() + pred_mask = torch.permute(pred_mask, (0, 2, 3, 1)) + pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)) masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1) poses = torch.concat((pose, pose_out, pred_pose), axis=2) @@ -121,7 +122,7 @@ if __name__ == '__main__': device = "cuda" if torch.cuda.is_available() else "cpu" - batch_size = 1 + batch_size = 4 use_cache = False train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache) @@ -156,8 +157,6 @@ if __name__ == '__main__': # if iteration == 0: # to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - if not args['debug']: - wandb.log(to_log, step=iteration) train_accumulator.reset() iteration += 1 # print(f'Iteration {iteration} finished') @@ -165,18 +164,17 @@ if __name__ == '__main__': epoch += 1 print(f'Epoch {epoch} finished') - if not iteration % 1000 == 0: - continue + # if not iteration % 1000 == 0: + # continue - for batch in valid_dataloader: - valid_loss = valid_step(batch) - valid_accumulator.update(valid_loss) + # for batch in valid_dataloader: + # valid_loss = valid_step(batch) + # valid_accumulator.update(valid_loss) - - to_log = dict({ **valid_loss, **{ 'current_step': iteration } }) - valid_figs = plot_everything(batch, 'train') - to_log = { **to_log, **valid_figs } - valid_accumulator.reset() + # to_log = dict({ **valid_loss, **{ 'current_step': iteration } }) + # valid_figs = plot_everything(batch, 'train') + # to_log = { **to_log, **valid_figs } + # valid_accumulator.reset() if not args['debug']: wandb.log(to_log, step=iteration)