use new colors for pose

data/pose_dataset.py

@@ -102,8 +102,8 @@ class Dataset(Dataset):
                     # 'body_out': body_out,
                     # 'render': render,
                     # 'render_out': render_out,
-                    'mask': mask,
-                    'mask_out': mask_out,
+                 'mask': mask,
+                 'mask_out': mask_out,
                }
 
         if self.cache is not None:

data/pose_renderer.py

@@ -30,9 +30,12 @@ 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()