pose pipeline trial

data/pose_dataset.py

@@ -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

@@ -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
+        mask, pose_in, pose_in_back, pose_out, mask_out = input
-        features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
+        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
+        return pose, mask_out_pred, rotation
-        x = pose
+        # 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 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

@@ -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)
+    masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy()
-    poses = torch.concat((pose, pose_out, pred_pose), axis=2)
+    poses = torch.concat((pose, pose_back, pose_out, pred_pose), axis=2).cpu().numpy()
-    masks_imgs = [wandb.Image((masks[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
+    masks_imgs = [wandb.Image((masks[i] + 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])]
+    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, } }
+
+    # 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
+# Use the profiler to track memory usage
-    epoch = 0
+    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
-    while True:
+    #     with record_function("model_inference"):
-
+    train()
-        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)
-
-
 
+    # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
 
     pass