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

pose pipeline trial

Browse Source
This commit is contained in:
luni
2023-06-12 13:59:53 +00:00
parent 33b4ccfcb4
commit 5f34b667fc
3 changed files with 95 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/pose_dataset.py
+11
View File
@@ -14,6 +14,7 @@ from cloth3d.read import DataReader
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
from superSMPL.smpl_torch import SMPL
from data.pose_renderer import PoseRenderer
import trimesh
class Dataset(Dataset):
@@ -76,13 +77,20 @@ class Dataset(Dataset):
body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
mesh = trimesh.Trimesh(vertices=body[0][0,:,:].cpu().numpy(), faces=self.smpl.f_all, process=False, validate=False)
mesh.apply_transform(trimesh.transformations.rotation_matrix(angle=3.14159, direction=[0, 1, 0], point=[0, 0, 0]))
body_back = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32)
body_back[0,:,0] = -body_back[0,:,0]
# print(f'smpl took {time() - start} seconds')
# start = time()
render = self._render_body(body[0])[0,:,:,:3]
render_back = self._render_body(body_back)[0,:,:,:3]
render_out = self._render_body(body_out[0])[0,:,:,:3]
render = render.permute(2,0,1)
render_back = render_back.permute(2,0,1)
render_out = render_out.permute(2,0,1)
# import plotly.express as ex
@@ -90,9 +98,11 @@ class Dataset(Dataset):
# ex.imshow(render_out[0,:,:,:3].cpu().numpy()).show()
mask = torch.min(render != 1, dim=0)[0][None,:,:].to(torch.float32)
mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32)
mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)
render = render * mask
render_back = render_back * mask_back
render_out = render_out * mask_out
# print(f'render took {time() - start} seconds')
@@ -101,6 +111,7 @@ class Dataset(Dataset):
'pose': render,
'shape': shape,
'pose_out': render_out,
'pose_back': render_back,
# 'body': body,
# 'body_out': body_out,
# 'render': render,
models/pose_encoder.py
+8 -7
View File
@@ -168,7 +168,7 @@ class PoseEncoder(nn.Module):
CNBlockConfig(16, None, 3),
]
self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
self.encoder = ConvNeXt(block_setting, input_features=7, device=device)
self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)
@@ -177,8 +177,8 @@ class PoseEncoder(nn.Module):
def forward(self, input) -> torch.Tensor:
mask, pose_in, pose_out = input
features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
mask, pose_in, pose_in_back, pose_out, mask_out = input
features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], dim=1))
x = features.view(features.size(0), features.size(1), -1)
y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
rotation = self.rotation_decoder(y)
@@ -188,14 +188,15 @@ class PoseEncoder(nn.Module):
# x = torch.einsum('bij,bkj->bik', x, rotation)
x = torch.einsum('bij,bjk->bik', x, rotation)
x = x.view(x.size(0), x.size(1), 16, 16)
pose, mask_out = self.decoder(x)
pose, mask_out_pred = self.decoder(x)
pose = pose * mask_out
# mask_out = torch.round(mask_out)
# mask_out = (mask_out > 0.5).to(torch.int64)
# pose = pose * (mask_out > 0.5).to(torch.float32)
# return pose, mask_out, rotation
x = pose
return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
return pose, mask_out_pred, rotation
# x = pose
# return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
class DumbPoseEncoder(nn.Module):
def __init__(self, in_channels=7, out_channels=3, device='cuda'):
scripts/pose_encoder_train.py
+76 -53
View File
@@ -26,6 +26,7 @@ import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.profiler import profile, record_function, ProfilerActivity
# from torchgeometry.losses.dice import DiceLoss
# from torchmetrics import Dice
@@ -40,7 +41,7 @@ ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
def init_wandb():
run = wandb.init(project='reconstruction', name='pose_1', notes='', dir='/data-local/127/temp', tags=['pose'])
run = wandb.init(project='reconstruction', name='pose_2', notes='have the back of the body also as the main input', dir='/data-local/127/temp', tags=['pose'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/pose_encoder_train.py')
artifact.add_file('models/pose_encoder.py')
@@ -69,30 +70,34 @@ class SymmetricLoss(nn.Module):
def train_step(batch):
model.train()
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
# pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
mse = loss_fn(pred_pose, batch['pose_out'])
mse_mask = 0 #loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
symm_rotation = 0 #loss_fn_rotation(pred_rotation)
loss = mse #+ mse_mask + symm_rotation
print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')
# print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.inference_mode():
mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
symm_rotation = loss_fn_rotation(pred_rotation)
return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/symm_rotation': symm_rotation }
def valid_step(batch):
model.eval()
with torch.inference_mode():
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
mse = loss_fn(pred_pose, batch['pose_out'])
mse_mask = loss_fn_mask(pred_mask, batch['mask_out'])
mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
symm_rotation = loss_fn_rotation(pred_rotation)
loss = mse + mse_mask + symm_rotation
loss = mse
return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
@@ -100,24 +105,76 @@ def valid_step(batch):
def plot_everything(batch, type: str):
model.eval()
with torch.inference_mode():
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_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_back = torch.permute(batch['pose_back'], (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))
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)
masks_imgs = [wandb.Image((masks[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
poses_imgs = [wandb.Image((poses[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } }
masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy()
poses = torch.concat((pose, pose_back, pose_out, pred_pose), axis=2).cpu().numpy()
masks_imgs = [wandb.Image((masks[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
poses_imgs = [wandb.Image((poses[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
# do a batchwise min-max normalization of pred_rotation
pred_rotation_img = pred_rotation.cpu().numpy()
pred_rotation_img = (pred_rotation_img - pred_rotation_img.min(axis=(1, 2,), keepdims=True)) / (pred_rotation_img.max(axis=(1, 2), keepdims=True) - pred_rotation_img.min(axis=(1, 2), keepdims=True))
pred_rotation_imgs = [wandb.Image(pred_rotation_img[i]) for i in range(pred_pose.shape[0])]
return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, f'{type}/rotation': pred_rotation_imgs} }
def train():
iteration = 0
epoch = 0
while True:
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if iteration % 1000 == 0 or (iteration < 500 and iteration % 100 == 0):
to_log = dict({ **train_loss, **{ 'current_step': iteration } })
if iteration % 5000 == 0 or (iteration < 5000 and iteration % 1000 == 0):
train_figs = plot_everything(batch, 'train')
to_log = { **to_log, **train_figs }
# if iteration == 0:
# to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
train_accumulator.reset()
# figure out when to log when using validation set as well
if not args['debug']:
wandb.log(to_log, step=iteration)
iteration += 1
# print(f'Iteration {iteration} finished')
epoch += 1
print(f'Epoch {epoch} finished')
# if not iteration % 1000 == 0:
# continue
# 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()
# if not args['debug']:
# wandb.log(to_log, step=iteration)
if __name__ == '__main__':
if not args['debug']:
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d')
init_wandb()
device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -141,45 +198,11 @@ if __name__ == '__main__':
lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
# optimizer = optim.SGD(model.parameters(), lr=0.1)
iteration = 0
epoch = 0
while True:
for batch in train_dataloader:
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if iteration % 1000 == 0 or (iteration < 1000 and iteration % 100 == 0):
to_log = dict({ **train_loss, **{ 'current_step': iteration } })
if iteration % 5000 == 0 or (iteration < 5000 and iteration % 1000 == 0):
train_figs = plot_everything(batch, 'train')
to_log = { **to_log, **train_figs }
# if iteration == 0:
# to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
train_accumulator.reset()
iteration += 1
# print(f'Iteration {iteration} finished')
epoch += 1
print(f'Epoch {epoch} finished')
# if not iteration % 1000 == 0:
# continue
# 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()
if not args['debug']:
wandb.log(to_log, step=iteration)
# Use the profiler to track memory usage
# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
# with record_function("model_inference"):
train()
# print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
pass