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

pose autoencoder pipeline

Browse Source
This commit is contained in:
luni
2023-06-16 22:11:57 +00:00
parent 5f34b667fc
commit 7a31f0d04c
4 changed files with 405 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/pose_dataset.py
+39 -22
View File
@@ -17,10 +17,11 @@ from data.pose_renderer import PoseRenderer
import trimesh
class Dataset(Dataset):
def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
class PoseDataset(Dataset):
def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True, ae_mode=False):
self.folder = folder
self.device = device
self.ae_mode = ae_mode
self.frames = read_frames_file(os.path.join(folder, frames_file))
self.reader = DataReader(folder, 'superSMPL')
@@ -49,7 +50,7 @@ class Dataset(Dataset):
# and rotate it 90 degress along the z axis
body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T
out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.get_device()))
out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.device))
return out
@@ -68,57 +69,73 @@ class Dataset(Dataset):
shape = torch.tensor(self.shapes[frame.id], device=self.device)
pose[:3] = 0
pose_out = pose.clone()
pose_out[1] = np.random.uniform(-np.pi, np.pi)
angle = torch.rand(1) * 2 * np.pi - np.pi
matrix = self.smpl.rodrigues(torch.tensor([0.0, angle.numpy()[0], 0.0], device=self.device)[None, None, :]).to(self.device)[0]
angle = angle.to(self.device)
pose_out[1] = angle
# pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
# generate random number between -np.pi, np.pi
# start = time()
body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
if self.ae_mode:
body = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
else:
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]
# 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_back = self._render_body(body_back)[0,:,:,:3]
if not self.ae_mode:
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)
# render_back = render_back.permute(2,0,1)
if not self.ae_mode:
render_out = render_out.permute(2,0,1)
# import plotly.express as ex
# ex.imshow(render[0,:,:,:3].cpu().numpy()).show()
# 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)
# mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32)
if not self.ae_mode:
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
# render_back = render_back * mask_back
if not self.ae_mode:
render_out = render_out * mask_out
# print(f'render took {time() - start} seconds')
data = {
'pose': render,
'shape': shape,
'pose_out': render_out,
'pose_back': render_back,
# 'pose_back': render_back,
# 'body': body,
# 'body_out': body_out,
# 'render': render,
# 'render_out': render_out,
'mask': mask,
'mask_out': mask_out,
'angle': angle,
'matrix': matrix,
}
if not self.ae_mode:
data['pose_out'] = render_out
data['mask_out'] = mask_out
if self.cache is not None:
self.cache[frame] = data
@@ -129,7 +146,7 @@ if __name__ == '__main__':
from torch.utils.data import DataLoader
device = "cuda" if torch.cuda.is_available() else "cpu"
dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
for batch in train_dataloader:
models/convnext.py
+1 -2
View File
@@ -159,11 +159,10 @@ class ConvNeXtBase(nn.Module):
def _forward_impl(self, x: Tensor) -> Tensor:
x = self.stem(x)
features: List[Tensor] = []
x = self.layers(x)
x = self.act(x)
return x, features
return x, None
def forward(self, x: Tensor) -> Tensor:
return self._forward_impl(x)
models/pose_ae.py
+183
View File
@@ -0,0 +1,183 @@
import os
import sys
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import torch
from torch import nn, Tensor
from torch.nn import functional as F
from functools import partial
from typing import Any, Callable, List, Optional, Sequence
torch.set_float32_matmul_precision('high')
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from models.convnext import ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
class ConvNeXtReversePose(nn.Module):
def __init__(
self,
block_setting: List[CNBlockConfig],
stochastic_depth_prob: float = 0.0,
layer_scale: float = 1e-6,
block: Optional[Callable[..., nn.Module]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None,
device: str = 'cuda',
use_skip=False,
) -> None:
super().__init__()
self.use_skip = use_skip
if not block_setting:
raise ValueError("The block_setting should not be empty")
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
raise TypeError("The block_setting should be List[CNBlockConfig]")
if block is None:
block = CNBlock
if norm_layer is None:
norm_layer = partial(LayerNorm2d, eps=1e-6)
stage: List[nn.Module] = []
for _ in range(3):
stage.append(block(block_setting[0].input_channels, layer_scale, 0))
self.stage_smooth = nn.Sequential(*stage).to(device)
self.first = nn.Sequential(
nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
norm_layer(block_setting[0].out_channels),
).to(device)
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0
layers_list: List[nn.Module] = []
for cnf in block_setting[1:]:
if cnf.input_channels is not None:
# Bottlenecks
stage: List[nn.Module] = []
for _ in range(cnf.num_layers):
# adjust stochastic depth probability based on the depth of the stage block
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
stage_block_id += 1
layers_list.append(nn.Sequential(*stage).to(device))
else:
layers_list.append(nn.Identity().to(device))
if cnf.out_channels is not None:
# Upsampling
layers_list.append(
nn.Sequential(
norm_layer(cnf.input_channels),
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
).to(device)
)
self.layers = nn.Sequential(*layers_list[:-1]).to(device)
self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
stage: List[nn.Module] = []
for _ in range(cnf.num_layers):
stage.append(block(block_setting[-1].input_channels, layer_scale, 0))
self.stage_mask = nn.Sequential(*stage).to(device)
self.last = nn.Conv2d(
block_setting[-1].input_channels,
3,
kernel_size=3,
padding='same',
).to(device)
self.last_mask = nn.Conv2d(
block_setting[-1].input_channels,
1,
kernel_size=3,
padding='same',
).to(device)
for m in self.modules():
if isinstance(m, (nn.Conv2d, nn.Linear)):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
def _forward_impl(self, x) -> Tensor:
if self.use_skip:
raise NotImplementedError
input, features = x
else:
input = x
x = self.stage_smooth(input)
x = self.first(x)
f = self.layers(x)
x = self.stage_pose(f)
pose = torch.sigmoid(self.last(x))
x = self.stage_mask(f)
mask = torch.sigmoid(self.last_mask(x))
pose = pose * (mask > 0.5).to(torch.float32)
return pose, mask
def forward(self, x) -> Tensor:
return self._forward_impl(x)
class PoseAE(nn.Module):
def __init__(self, features=256, device='cuda'):
super().__init__()
block_setting = [
CNBlockConfig(96, 192, 3),
CNBlockConfig(192, 384, 3),
CNBlockConfig(384, features, 9),
CNBlockConfig(features, None, 3),
]
block_setting_reverse = [
CNBlockConfig(features, 384, 3),
CNBlockConfig(384, 192, 9),
CNBlockConfig(192, 96, 3),
CNBlockConfig(96, 27, 3),
CNBlockConfig(27, 9, 3),
CNBlockConfig(9, None, 3),
]
self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
self.decoder = ConvNeXtReversePose(block_setting_reverse, device=device)
self.flatten = nn.Flatten()
def forward(self, input) -> torch.Tensor:
mask, pose_in = input
x, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
pose_out, mask_out = self.decoder(x)
return pose_out, mask_out
if __name__ == '__main__':
# test the model
model = PoseAE()
# model = torch.compile(model)
print(model)
print(sum(p.numel() for p in model.parameters() if p.requires_grad))
print(sum(p.numel() for p in model.parameters()))
x = torch.rand((4, 1, 512, 512)).to('cuda')
y = torch.rand((4, 3, 512, 512)).to('cuda')
output = model((x, y))
pass
scripts/pose_ae_train.py
+182
View File
@@ -0,0 +1,182 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import argparse
import importlib
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
# config_module = importlib.import_module(args['config'])
# config = config_module.config
if args['debug']:
pass
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '7'
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import wandb
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
from data.pose_dataset import PoseDataset
from utils.utils import Accumulator
from models.pose_ae import PoseAE
from models.convnext import ConvNeXt
import plotly.express as ex
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
def init_wandb():
run = wandb.init(project='reconstruction', name='pose_ae_2', notes='now with random poses', dir='/data-local/127/temp', tags=['pose'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/pose_ae_train.py')
artifact.add_file('models/pose_ae.py')
artifact.add_file('data/pose_dataset.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
def train_step(batch):
model.train()
pred_pose, pred_mask = model([batch['mask'], batch['pose']])
mse = loss_fn(pred_pose, batch['pose'])
mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask'].squeeze(1))
loss: torch.Tensor = mse + mse_mask
optimizer.zero_grad()
loss.backward()
optimizer.step()
return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask }
def valid_step(batch):
model.eval()
with torch.inference_mode():
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.squeeze(1), batch['mask_out'].squeeze(1))
symm_rotation = loss_fn_rotation(pred_rotation)
loss = mse
return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
def plot_everything(batch, type: str):
model.eval()
with torch.inference_mode():
pred_pose, pred_mask = model([batch['mask'], batch['pose']])
mask = torch.permute(batch['mask'], (0, 2, 3, 1))
pose = torch.permute(batch['pose'], (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, pred_mask), axis=2).squeeze(-1).cpu().numpy()
poses = torch.concat((pose, 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])]
return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_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"
batch_size = 4
use_cache = False
train_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=True)
valid_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=True)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
model = PoseAE(device=device).to(device)
loss_fn = nn.MSELoss()
loss_fn_mask = nn.BCELoss(reduction='mean')
optimizer = optim.Adam(model.parameters(), lr=0.001)
# lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
# optimizer = optim.SGD(model.parameters(), lr=0.1)
# 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