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

use new colors for pose

Browse Source
This commit is contained in:
luni
2023-06-08 10:54:34 +00:00
parent fdc876025a
commit 882ddf5578
3 changed files with 23 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/pose_dataset.py
+2 -2
View File
@@ -102,8 +102,8 @@ class Dataset(Dataset):
# 'body_out': body_out, # 'body_out': body_out,
# 'render': render, # 'render': render,
# 'render_out': render_out, # 'render_out': render_out,
'mask': mask, 'mask': mask,
'mask_out': mask_out, 'mask_out': mask_out,
} }
if self.cache is not None: if self.cache is not None:
data/pose_renderer.py
+6 -3
View File
@@ -30,9 +30,12 @@ from pytorch3d.structures import Meshes
class PoseRenderer(Module): class PoseRenderer(Module):
def __init__(self, device='cuda') -> None: def __init__(self, device='cuda') -> None:
super().__init__() super().__init__()
colors = np.load('custom_data/pose_colors.npy') # 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/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.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() self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
scripts/pose_encoder_train.py
+15 -17
View File
@@ -94,19 +94,20 @@ def valid_step(batch):
symm_rotation = loss_fn_rotation(pred_rotation) symm_rotation = loss_fn_rotation(pred_rotation)
loss = mse + mse_mask + symm_rotation loss = mse + mse_mask + symm_rotation
return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': 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): 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 = torch.permute(batch['mask'], (0, 2, 3, 1))
mask_out = torch.permute(batch['mask_out'], (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 = torch.permute(batch['pose'], (0, 2, 3, 1))
pose_out = torch.permute(batch['pose_out'], (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_mask = torch.permute(pred_mask, (0, 2, 3, 1))
pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)).detach() pred_pose = torch.permute(pred_pose, (0, 2, 3, 1))
masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1) masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1)
poses = torch.concat((pose, pose_out, pred_pose), axis=2) 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" device = "cuda" if torch.cuda.is_available() else "cpu"
batch_size = 1 batch_size = 4
use_cache = False use_cache = False
train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache) 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: # if iteration == 0:
# to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } # 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() train_accumulator.reset()
iteration += 1 iteration += 1
# print(f'Iteration {iteration} finished') # print(f'Iteration {iteration} finished')
@@ -165,18 +164,17 @@ if __name__ == '__main__':
epoch += 1 epoch += 1
print(f'Epoch {epoch} finished') print(f'Epoch {epoch} finished')
if not iteration % 1000 == 0: # if not iteration % 1000 == 0:
continue # continue
for batch in valid_dataloader: # for batch in valid_dataloader:
valid_loss = valid_step(batch) # valid_loss = valid_step(batch)
valid_accumulator.update(valid_loss) # valid_accumulator.update(valid_loss)
# to_log = dict({ **valid_loss, **{ 'current_step': iteration } })
to_log = dict({ **valid_loss, **{ 'current_step': iteration } }) # valid_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(batch, 'train') # to_log = { **to_log, **valid_figs }
to_log = { **to_log, **valid_figs } # valid_accumulator.reset()
valid_accumulator.reset()
if not args['debug']: if not args['debug']:
wandb.log(to_log, step=iteration) wandb.log(to_log, step=iteration)