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,