use new colors for pose

data/pose_renderer.py

@@ -30,8 +30,11 @@ 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()

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):
+    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)